#include "graph_creator.h"

#include <library/cpp/testing/gtest/gtest.h>
#include <maps/wikimap/mapspro/services/mrc/long_tasks/tasks_gen/lib/minimal_overhead.h>

namespace maps::mrc::gen_targets::tests {

// Checks that addedEdge from ExtendedRoadNetwork
// is a copy of originalEdge from orignal RoadNetword.
// Checks all the edge data except isTarget flag.
bool edgesAreEqual(const ExtendedRoadNetwork& roadNetwork,
                   const Edge& originalEdge, const Edge& addedEdge) {
    // filter out temporary added deadend edges
    std::vector<EdgeId> addedEdgeInEdges;
    std::vector<EdgeId> addedEdgeOutEdges;
    std::vector<Seconds> addedEdgeOutPenalties;
    for (size_t i = 0; i < addedEdge.outEdges.size(); i++) {
        const Edge& outEdge = roadNetwork.edge(addedEdge.outEdges[i]);
        if (outEdge.outEdges.size() && outEdge.inEdges.size()) {
            addedEdgeOutEdges.push_back(addedEdge.outEdges[i]);
            addedEdgeOutPenalties.push_back(addedEdge.outEdgesPenalties[i]);
        }
    }
    for (size_t i = 0; i < addedEdge.inEdges.size(); i++) {
         const Edge& inEdge = roadNetwork.edge(addedEdge.inEdges[i]);
        if (inEdge.outEdges.size() && inEdge.inEdges.size()) {
            addedEdgeInEdges.push_back(addedEdge.inEdges[i]);
        }
    }

    return (originalEdge.geom.points() == addedEdge.geom.points() &&
            originalEdge.fc == addedEdge.fc &&
            originalEdge.length == addedEdge.length &&
            originalEdge.time == addedEdge.time &&
            originalEdge.isUTurn == addedEdge.isUTurn &&
            originalEdge.inEdges == addedEdgeInEdges &&
            originalEdge.outEdges == addedEdgeOutEdges &&
            originalEdge.outEdgesPenalties == addedEdgeOutPenalties);
}

TEST(graph_optimization_tests, test_extended_road_network_nodes)
{
    //  graph:
    //     |
    //     5
    //     v
    //--1-->==2===>--3-->--4--->
    //     |
    //     6
    //     v
    //
    // edge 2 is target
    // there should be created 2 nodes near edge2

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2, 6}, false)
        .addEdge(2, {20, 0}, {40, 0}, {3})
        .addEdge(3, {60, 0}, {80, 0}, {4}, false)
        .addEdge(4, {80, 0}, {100, 0}, {}, false)
        .addEdge(5, {80, 0}, {100, 0}, {6}, false)
        .addEdge(6, {80, 0}, {100, 0}, {}, false);
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());

    std::vector<Node> nodes;
    for (const auto& nodeIt : roadNetwork.nodes()) {
        nodes.push_back(nodeIt.second);
    }
    EXPECT_EQ(nodes.size(), 2u);

    if (nodes[0].outEdges.size() == 1 && nodes[0].outEdges[0] == 3) {
        EXPECT_THAT(nodes[0].inEdges, testing::UnorderedElementsAre(2));
        EXPECT_THAT(nodes[1].outEdges, testing::UnorderedElementsAre(2, 6));
        EXPECT_THAT(nodes[1].inEdges, testing::UnorderedElementsAre(1, 5));
    } else {
        EXPECT_THAT(nodes[0].outEdges, testing::UnorderedElementsAre(2, 6));
        EXPECT_THAT(nodes[0].inEdges, testing::UnorderedElementsAre(1, 5));
        EXPECT_THAT(nodes[1].outEdges, testing::UnorderedElementsAre(3));
        EXPECT_THAT(nodes[1].inEdges, testing::UnorderedElementsAre(2));
    }
}

TEST(graph_optimization_tests, test_extended_road_network_generate_id)
{
    //  graph:
    //
    //==1===>
    //<=2====
    //

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {3})
        .addEdge(3, {20, 0}, {0, 0}, {1});
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());

    EdgeId newEdgeId = roadNetwork.generateNewEdgeId();
    EXPECT_GT(newEdgeId, 3);
    EXPECT_GT(roadNetwork.generateNewEdgeId(), newEdgeId);
}

TEST(graph_optimization_tests, test_extended_road_network_add_edge)
{
    //  graph:
    //
    //==1===>--4-->--3-->--5--->
    //<=2====
    //
    // adding edge 4

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2})
        .addEdge(2, {20, 0}, {0, 0}, {1})
        .addEdge(3, {40, 0}, {60, 0}, {5}, false)
        .addEdge(5, {60, 0}, {80, 0}, {}, false);
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());

    GraphCreator graphCreator2;
    graphCreator
        .addEdge(4, {20, 0}, {40, 0}, {}, false);
    Edge edge4 = graphCreator.createGraph().edge(4);
    edge4.inEdges = {1};
    edge4.outEdges = {3};
    edge4.outEdgesPenalties = {10};

    roadNetwork.addEdge(edge4);

    EXPECT_TRUE(roadNetwork.getEdges().count(4));
    const Edge& addedEdge = roadNetwork.edge(4);
    EXPECT_TRUE(edgesAreEqual(roadNetwork, edge4, addedEdge));
    EXPECT_THAT(roadNetwork.edge(1).outEdges, testing::UnorderedElementsAre(2, 4));
    EXPECT_THAT(roadNetwork.edge(3).inEdges, testing::UnorderedElementsAre(4));

    std::vector<Node> nodes;
    for (const auto& nodeIt : roadNetwork.nodes()) {
        nodes.push_back(nodeIt.second);
    }
    EXPECT_EQ(nodes.size(), 2u);

    if (nodes[0].outEdges.size() == 1 && nodes[0].outEdges[0] == 1) {
        EXPECT_THAT(nodes[0].inEdges, testing::UnorderedElementsAre(2));
        EXPECT_THAT(nodes[1].outEdges, testing::UnorderedElementsAre(2, 4));
        EXPECT_THAT(nodes[1].inEdges, testing::UnorderedElementsAre(1));
    } else {
        EXPECT_THAT(nodes[0].outEdges, testing::UnorderedElementsAre(2, 4));
        EXPECT_THAT(nodes[0].inEdges, testing::UnorderedElementsAre(1));
        EXPECT_THAT(nodes[1].outEdges, testing::UnorderedElementsAre(1));
        EXPECT_THAT(nodes[1].inEdges, testing::UnorderedElementsAre(2));
    }
}

TEST(graph_optimization_tests, test_extended_road_network_connect_edges)
{
    //  graph:
    //
    //--1-->==2===>--3-->
    //
    //

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2}, false)
        .addEdge(2, {20, 0}, {40, 0}, {3})
        .addEdge(3, {40, 0}, {60, 0}, {});
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());

    EXPECT_THAT(roadNetwork.edge(1).inEdges, testing::UnorderedElementsAre());
    EXPECT_THAT(roadNetwork.edge(3).outEdges, testing::UnorderedElementsAre());
    EXPECT_EQ(roadNetwork.edge(3).outEdgesPenalties.size(), 0u);

    roadNetwork.connectEdges(3, 1, 2.0);
    EXPECT_THAT(roadNetwork.edge(1).inEdges, testing::UnorderedElementsAre(3));
    EXPECT_THAT(roadNetwork.edge(3).outEdges, testing::UnorderedElementsAre(1));
    EXPECT_EQ(roadNetwork.edge(3).outEdgesPenalties.size(), 1u);
    EXPECT_EQ(roadNetwork.edge(1).time, 2.0);

    roadNetwork.disconnectEdges(3, 1);
    EXPECT_THAT(roadNetwork.edge(1).inEdges, testing::UnorderedElementsAre());
    EXPECT_THAT(roadNetwork.edge(3).outEdges, testing::UnorderedElementsAre());
    EXPECT_EQ(roadNetwork.edge(3).outEdgesPenalties.size(), 0u);
}

TEST(graph_optimization_tests, test_zero_optimal_overhead)
{
    //  graph:
    //
    //==1===>
    //<=2====
    //

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2})
        .addEdge(2, {20, 0}, {0, 0}, {1});
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());
    ExtendedRoadNetwork optimizedRoadNetwork = roadNetwork;

    auto optimalOverheadCreator = OptimalOverheadCreator(optimizedRoadNetwork);
    // no overhead edges should be added
    EXPECT_EQ(optimizedRoadNetwork.getEdges().size(), 2u);
}

TEST(graph_optimization_tests, test_zero_optimal_overhead2)
{
    //  graph:
    //
    //==1==>--3-->
    //<==2==<--4--
    //
    // edges 1, 2 are targets

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2})
        .addEdge(2, {20, 0}, {0, 0}, {1})
        .addEdge(3, {20, 0}, {40, 0}, {2}, false)
        .addEdge(4, {40, 0}, {20, 0}, {1}, false);

    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());
    ExtendedRoadNetwork optimizedRoadNetwork = roadNetwork;

    auto optimalOverheadCreator = OptimalOverheadCreator(optimizedRoadNetwork);
    // no overhead edges should be added
    EXPECT_EQ(optimizedRoadNetwork.getEdges().size(), 4u);
}

TEST(graph_optimization_tests, test_simple_optimal_overhead)
{
    //  graph:
    //
    //==1==>
    //<--2--
    //
    // edge 1 is target
    // optimal overhead is edge 2

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2})
        .addEdge(2, {20, 0}, {0, 0}, {1}, false);
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());
    ExtendedRoadNetwork optimizedRoadNetwork = roadNetwork;

    auto optimalOverheadCreator = OptimalOverheadCreator(optimizedRoadNetwork);

    // checks added adges
    std::vector<EdgeId> addedEdges;
    for (auto& edge : optimizedRoadNetwork.edges()) {
        // select only added edges
        if (edge.first == 1 || edge.first == 2) {
            continue;
        }
        if (!edge.second.outEdges.size() || !edge.second.inEdges.size()) {
            // edge is a temporary deadend edge
            continue;
        }

        // optimizer should add only target edges
        EXPECT_TRUE(edge.second.isTarget);
        EdgeId originalId = optimalOverheadCreator.getOriginalId(edge.first);
        addedEdges.push_back(originalId);
        EXPECT_TRUE(edgesAreEqual(optimizedRoadNetwork, roadNetwork.edge(2), edge.second));

    }

    EXPECT_THAT(addedEdges, testing::UnorderedElementsAre(2));
}

TEST(graph_optimization_tests, test_optimal_overhead)
{
    //  graph:
    //
    //==1===> --2-> ===3==> --4-> ==5===>
    //<-6---- <-7-- <---8-- <-9-- <--10--
    //
    // edges 1, 3, 5 are targets
    // optimal overhead - edges 6, 8, 10

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {20, 0}, {2, 6})
        .addEdge(2, {20, 0}, {30, 0}, {3, 7}, false)
        .addEdge(3, {30, 0}, {50, 0}, {4, 8})
        .addEdge(4, {50, 0}, {60, 0}, {5, 9}, false)
        .addEdge(5, {60, 0}, {80, 0}, {10})
        .addEdge(6, {20, 0}, {0, 0}, {1}, false)
        .addEdge(7, {30, 0}, {20, 0}, {2, 6}, false)
        .addEdge(8, {50, 0}, {30, 0}, {3, 7}, false)
        .addEdge(9, {60, 0}, {50, 0}, {4, 8}, false)
        .addEdge(10, {80, 0}, {60, 0}, {5, 9}, false);
    ExtendedRoadNetwork roadNetwork(graphCreator.createGraph().getEdges());
    ExtendedRoadNetwork optimizedRoadNetwork = roadNetwork;

    auto optimalOverheadCreator = OptimalOverheadCreator(optimizedRoadNetwork);

    std::set<EdgeId> originalEdges;
    for (auto& edge : roadNetwork.edges()) {
        originalEdges.insert(edge.first);
    }

    // checks added adges
    std::vector<EdgeId> addedEdges;
    for (auto& edge : optimizedRoadNetwork.edges()) {
        // select only added edges
        if (originalEdges.count(edge.first)) {
            continue;
        }
        if (!edge.second.outEdges.size() || !edge.second.inEdges.size()) {
            // edge is a temporary deadend edge
            continue;
        }

        // optimizer should add only target edges
        EXPECT_TRUE(edge.second.isTarget);
        EdgeId originalId = optimalOverheadCreator.getOriginalId(edge.first);
        addedEdges.push_back(originalId);
        if (originalId == 6) {
            EXPECT_TRUE(edgesAreEqual(optimizedRoadNetwork, roadNetwork.edge(6), edge.second));
        } else if(originalId == 8) {
            EXPECT_TRUE(edgesAreEqual(optimizedRoadNetwork, roadNetwork.edge(8), edge.second));
        } else {
            EXPECT_TRUE(edgesAreEqual(optimizedRoadNetwork, roadNetwork.edge(10), edge.second));
        }

    }

    EXPECT_THAT(addedEdges, testing::UnorderedElementsAre(6, 8, 10));
}

} // namespace maps::mrc::gen_targets::tests
