#include "transport_renderer.h"

#include <algorithm>
#include <sstream>
#include <limits>
#include <set>
#include <cassert>

TransportRenderer::TransportRenderer(const Json::Dict& render_settings_json) {
    mSettings = RenderSettings::fromJson(render_settings_json);
    SetupLayerBuilders();
}

void TransportRenderer::Init(std::vector<Stop> stops, std::vector<Bus> buses) {
    SetStops(std::move(stops));
    SetBuses(std::move(buses));
}

void TransportRenderer::SetStops(std::vector<Descriptions::Stop> stops) {
    mStops = std::move(stops);

    std::sort(mStops.begin(), mStops.end(), [](const Stop& a, const Stop& b) {
        return a.name < b.name;
    });

    for (auto& stop : mStops) {
        mNameToStop.insert({stop.name, &stop});
    }

}

void TransportRenderer::SetBuses(std::vector<Descriptions::Bus> buses) {
    mBuses = std::move(buses);
    std::sort(mBuses.begin(), mBuses.end(), [](const Bus& a, const Bus& b) {
        return a.name < b.name;
    });

    InterpolateInternalStops();
    ComputeLatLonIndices();
    ComputeRenderingSteps();
}

std::string TransportRenderer::Render() const {
    Svg::Document doc;

    for (const auto& layer : mSettings.layers) {
        mBuilders.at(layer)(doc);
    }

    std::stringstream sstream;
    doc.Render(sstream);

#ifdef DEBUG
    return sstream.str();
#else
    auto rawString = sstream.str();
    std::string convertedString;
    for (auto ch : rawString) {
        if (ch == '\\' || ch == '"')
            convertedString += '\\';
        convertedString += ch;
    }
    return convertedString;
#endif
}

Svg::Color ColorFromJson(const Json::Node& node) {
    if (std::holds_alternative<std::string>(node.GetBase())) {
        return Svg::Color(node.AsString());
    } else {
        if (node.AsArray().size() == 3) {
            return Svg::Color(Svg::Rgb{node.AsArray()[0].AsInt(),
                                       node.AsArray()[1].AsInt(),
                                       node.AsArray()[2].AsInt()
            });
        } else {
            return Svg::Color(Svg::Rgba{node.AsArray()[0].AsInt(),
                                        node.AsArray()[1].AsInt(),
                                        node.AsArray()[2].AsInt(),
                                        node.AsArray()[3].AsDouble()
            });
        }
    }
}

TransportRenderer::RenderSettings TransportRenderer::RenderSettings::fromJson(const Json::Dict& render_settings_json) {
    TransportRenderer::RenderSettings settings;

    settings.width = render_settings_json.at("width").AsDouble();
    settings.height = render_settings_json.at("height").AsDouble();
    settings.padding = render_settings_json.at("padding").AsDouble();
    settings.stopRadius = render_settings_json.at("stop_radius").AsDouble();
    settings.lineWidth = render_settings_json.at("line_width").AsDouble();
    settings.stopLabelFontSize = render_settings_json.at("stop_label_font_size").AsInt();
    settings.stopLabelOffset = {render_settings_json.at("stop_label_offset").AsArray()[0].AsDouble(),
                                render_settings_json.at("stop_label_offset").AsArray()[1].AsDouble()};
    settings.underlayerWidth = render_settings_json.at("underlayer_width").AsDouble();
    settings.underlayerColor = ColorFromJson(render_settings_json.at("underlayer_color"));
    for (const auto& colorNode : render_settings_json.at("color_palette").AsArray()) {
        settings.colorPalette.push_back(ColorFromJson(colorNode));
    }
    settings.busLabelFontSize = render_settings_json.at("bus_label_font_size").AsInt();
    settings.busLabelOffset = {render_settings_json.at("bus_label_offset").AsArray()[0].AsDouble(),
                               render_settings_json.at("bus_label_offset").AsArray()[1].AsDouble()};

    for (const auto& layerNode : render_settings_json.at("layers").AsArray()) {
        settings.layers.push_back(layerNode.AsString());
    }

    return settings;
}

void TransportRenderer::SetupLayerBuilders() {
    mBuilders["bus_lines"] = [this](Svg::Document& document) {
        for (size_t busIndex = 0; busIndex < mBuses.size(); ++busIndex) {
            auto color = mSettings.colorPalette[busIndex % mSettings.colorPalette.size()];

            auto polyline = Svg::Polyline{}.SetStrokeColor(color)
                                           .SetStrokeWidth(mSettings.lineWidth)
                                           .SetStrokeLineCap("round")
                                           .SetStrokeLineJoin("round");
            for (const auto& stopName : mBuses[busIndex].stops) {
                const auto& stop = *mNameToStop.at(stopName);
                polyline = polyline.AddPoint(GlobalToLocal(stop));
            }

            document.Add(polyline);
        }
    };

    mBuilders["bus_labels"] = [this](Svg::Document& document) {
        auto AddBusStopName = [this, &document](const Bus& bus, const Stop& stop, const Svg::Color color) {
            document.Add(Svg::Text{}.SetPoint(GlobalToLocal(stop))
                                    .SetOffset(mSettings.busLabelOffset)
                                    .SetFillColor(color)
                                    .SetFontSize(mSettings.busLabelFontSize)
                                    .SetFontFamily("Verdana")
                                    .SetFontWeight("bold")
                                    .SetData(bus.name)
                                    .SetFillColor(mSettings.underlayerColor)
                                    .SetStrokeColor(mSettings.underlayerColor)
                                    .SetStrokeWidth(mSettings.underlayerWidth)
                                    .SetStrokeLineCap("round")
                                    .SetStrokeLineJoin("round"));

            document.Add(Svg::Text{}.SetPoint(GlobalToLocal(stop))
                                    .SetOffset(mSettings.busLabelOffset)
                                    .SetFillColor(color)
                                    .SetFontSize(mSettings.busLabelFontSize)
                                    .SetFontFamily("Verdana")
                                    .SetFontWeight("bold")
                                    .SetData(bus.name));
        };

        for (size_t busIndex = 0; busIndex < mBuses.size(); ++busIndex) {
            auto color = mSettings.colorPalette[busIndex % mSettings.colorPalette.size()];
            const auto& bus = mBuses[busIndex];
            auto stopName = bus.stops[0];
            AddBusStopName(mBuses[busIndex], *mNameToStop.at(stopName), color);
            if (!bus.isRoundtrip && stopName != bus.stops[bus.stops.size() / 2])
                AddBusStopName(mBuses[busIndex], *mNameToStop.at(bus.stops[bus.stops.size() / 2]), color);
        }
    };

    mBuilders["stop_points"] = [this](Svg::Document& document) {
        for (const auto& stop : mStops) {
            document.Add(Svg::Circle{}.SetCenter(GlobalToLocal(stop))
                                      .SetRadius(mSettings.stopRadius)
                                      .SetFillColor("white"));
        }
    };

    mBuilders["stop_labels"] = [this](Svg::Document& doc) {
        for (const auto& stop : mStops) {
            doc.Add(Svg::Text{}.SetPoint(GlobalToLocal(stop))
                               .SetOffset(mSettings.stopLabelOffset)
                               .SetFontSize(mSettings.stopLabelFontSize)
                               .SetFontFamily("Verdana")
                               .SetData(stop.name)
                               .SetFillColor(mSettings.underlayerColor)
                               .SetStrokeColor(mSettings.underlayerColor)
                               .SetStrokeWidth(mSettings.underlayerWidth)
                               .SetStrokeLineCap("round")
                               .SetStrokeLineJoin("round"));

            doc.Add(Svg::Text{}.SetPoint(GlobalToLocal(stop))
                               .SetOffset(mSettings.stopLabelOffset)
                               .SetFontSize(mSettings.stopLabelFontSize)
                               .SetFontFamily("Verdana")
                               .SetData(stop.name)
                               .SetFillColor("black"));

        }
    };
}

void TransportRenderer::InterpolateInternalStops() {
    std::set<std::string> pivotStops;

    std::map<std::string, int> stopToBuses;
    for (const auto& bus : mBuses) {
        if (bus.stops.empty())
            continue;

        pivotStops.insert(bus.stops.front());
        if (!bus.isRoundtrip) {
            pivotStops.insert(bus.stops[bus.stops.size() / 2]);
        }
        std::map<std::string, int> stopToCount;
        for (const auto& stop : bus.stops) {
            stopToCount[stop]++;
        }

        for (const auto& [stop, count] : stopToCount) {
            if (count > 2)
                pivotStops.insert(stop);
            stopToBuses[stop]++;
        }
    }

    for (const auto& [stop, count] : stopToBuses) {
        if (count > 1) {
            pivotStops.insert(stop);
        }
    }

    for (const auto& bus : mBuses) {
        if (bus.stops.size() == 0)
            continue;

        int lastPivotIndex = 0;
        auto lastPivotStop = bus.stops[0];
        for (size_t i = 0; i < bus.stops.size(); ++i) {
            auto curStop = bus.stops[i % bus.stops.size()];
            if (pivotStops.count(curStop) > 0) {
                auto startPosition = mNameToStop[lastPivotStop]->position;
                auto endPosition = mNameToStop[curStop]->position;
                for (int j = lastPivotIndex + 1; j < i; ++j) {
                    auto& stop = *mNameToStop[bus.stops[j]];
                    assert(pivotStops.count(stop.name) == 0);
                    stop.position.latitude = startPosition.latitude + (j - lastPivotIndex) *
                                             (endPosition.latitude - startPosition.latitude) / (i - lastPivotIndex);
                    stop.position.longitude = startPosition.longitude + (j - lastPivotIndex) *
                                              (endPosition.longitude - startPosition.longitude) / (i - lastPivotIndex);
                }
                lastPivotIndex = static_cast<int>(i);
                lastPivotStop = curStop;
            }
        }
    }

    mPivots = pivotStops;
}

void TransportRenderer::ComputeLatLonIndices() {
    auto getSortedCoordinates = [this](auto coordGetter) {
        std::vector<std::pair<double, Stop>> result;
        for (const auto& stop : mStops) {
            result.push_back({coordGetter(stop), stop});
        }
        std::sort(result.begin(), result.end(), [](const std::pair<double, Stop>& a, const std::pair<double, Stop>& b) {
            return a.first + 1e-6 < b.first;
        });
        return result;
    };

    const std::vector<std::pair<double, Stop>> sortedLats = getSortedCoordinates([](const Stop& stop) {
        return stop.position.latitude;
    });
    const std::vector<std::pair<double, Stop>> sortedLons = getSortedCoordinates([](const Stop& stop) {
        return stop.position.longitude;
    });

    auto computePositionIndices = [this](std::map<std::string, int>& stopToIndex, const std::vector<std::pair<double, Stop>>& sortedCoordinate) {
        std::vector<Stop> currentStops {sortedCoordinate[0].second};
        int curId = 0;
        stopToIndex[currentStops.back().name] = curId;

        size_t startIndex = 1;
        for (; startIndex < sortedCoordinate.size(); ++startIndex) {
            if (std::abs(sortedCoordinate[startIndex].first - sortedCoordinate[0].first) < 1e-6) {
                currentStops.push_back(sortedCoordinate[startIndex].second);
                stopToIndex[currentStops.back().name] = curId;
            } else {
                break;
            }
        }

        while (startIndex < sortedCoordinate.size()) {
            size_t finishIndex = startIndex;
            while (finishIndex < sortedCoordinate.size() && std::abs(sortedCoordinate[startIndex].first - sortedCoordinate[finishIndex].first) < 1e-6) {
                finishIndex++;
            }

            bool shouldUnion = true;

            for (size_t i = startIndex; i < finishIndex; ++i) {
                Stop stop = sortedCoordinate[i].second;
                for (const auto& bus : mBuses) {
                    for (int j = 0; j < static_cast<int>(bus.stops.size()); ++j) {
                        if (bus.stops[j] == stop.name) {
                            auto busStop = *mNameToStop[bus.stops[(j + 1) % bus.stops.size()]];
                            if (std::find(currentStops.begin(), currentStops.end(), busStop) != currentStops.end()) {
                                shouldUnion = false;
                            }
                        }
                        if (bus.stops[(j + 1) % bus.stops.size()] == stop.name) {
                            auto busStop = *mNameToStop[bus.stops[j]];
                            if (std::find(currentStops.begin(), currentStops.end(), busStop) != currentStops.end()) {
                                shouldUnion = false;
                            }
                        }
                    }
                }
                currentStops.push_back(stop);
            }

            for (size_t i = startIndex; i < finishIndex; ++i) {
                currentStops.pop_back();
            }

            if (!shouldUnion) {
                curId++;
                currentStops.clear();
            }
            for (size_t i = startIndex; i < finishIndex; ++i) {
                currentStops.push_back(sortedCoordinate[i].second);
                stopToIndex[currentStops.back().name] = curId;
            }

            startIndex = finishIndex;
        }
    };

    computePositionIndices(mLatIndex, sortedLats);
    computePositionIndices(mLonIndex, sortedLons);
}

void TransportRenderer::ComputeRenderingSteps() {
    int maxLonIndex = 0;
    for (const auto& [key, value] : mLonIndex) {
        maxLonIndex = std::max(maxLonIndex, value);
    }

    int maxLatIndex = 0;
    for (const auto& [key, value] : mLatIndex) {
        maxLatIndex = std::max(maxLatIndex, value);
    }

    mLongitudeStep = (maxLonIndex == 0 ? 0 : (mSettings.width - 2 * mSettings.padding) / maxLonIndex);
    mLatitudeStep = (maxLatIndex == 0 ? 0 : (mSettings.height - 2 * mSettings.padding) / maxLatIndex);
}

Svg::Point TransportRenderer::GlobalToLocal(const Stop& stop) const {
    double x = mSettings.padding + mLonIndex.at(stop.name) * mLongitudeStep;
    double y = mSettings.height - mSettings.padding - mLatIndex.at(stop.name) * mLatitudeStep;

//    std::cout << "Stop: " << stop.name << ". " << x << " " << y << std::endl;

    return {x, y};
}