#include "graph_creator.h"

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

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

TEST(connected_components_tests, test_one_component)
{
    //  graph:
    //     --5->
    //     ^    |
    //     |    6
    //     4    |
    //     |    V
    //--1-> <-2-- --3->

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {1, 0}, {4})
        .addEdge(2, {2, 0}, {1, 0}, {4})
        .addEdge(3, {2, 0}, {3, 0}, {})
        .addEdge(4, {1, 0}, {1, 1}, {5})
        .addEdge(5, {1, 1}, {2, 1}, {6})
        .addEdge(6, {2, 1}, {2, 0}, {2, 3});

    RoadNetworkData roadNetwork = graphCreator.createGraph();
    {
        auto edgeIdToComponentId = getStronglyConnectedComponents(roadNetwork);
        EXPECT_EQ(edgeIdToComponentId.size(), 6u);
        EXPECT_EQ(edgeIdToComponentId.at(2), edgeIdToComponentId.at(4));
        EXPECT_EQ(edgeIdToComponentId.at(2), edgeIdToComponentId.at(5));
        EXPECT_EQ(edgeIdToComponentId.at(2), edgeIdToComponentId.at(6));
        EXPECT_NE(edgeIdToComponentId.at(2), edgeIdToComponentId.at(1));
        EXPECT_NE(edgeIdToComponentId.at(2), edgeIdToComponentId.at(3));
        EXPECT_NE(edgeIdToComponentId.at(1), edgeIdToComponentId.at(3));
    }

    {
        auto components = getFilteredConnectedComponentsOfTargets(
            roadNetwork, {1, 2, 3, 4, 5, 6});
        EXPECT_EQ(components.size(), 1u);
        std::set<EdgeId> result(components[0].begin(), components[0].end());
        EXPECT_THAT(result, testing::UnorderedElementsAre(2, 4, 5, 6));
    }

    {
        // if edge 6 is not target, component will still be connected
        auto components = getFilteredConnectedComponentsOfTargets(
            roadNetwork, {2, 4, 5});
        EXPECT_EQ(components.size(), 1u);
        std::set<EdgeId> result(components[0].begin(), components[0].end());
        EXPECT_THAT(result, testing::UnorderedElementsAre(2, 4, 5));
    }
}

TEST(connected_components_tests, test_two_disconnected_components)
{
    //  graph:
    // --5->       --9-->
    // ^    |      ^    |
    // |    6       \   10
    // 4    |        8  |
    // |    v         \ v
    // <-2--

    // Take only the largest strongly connected component
    // We consider that all the small components
    // are closed areas

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {1, 0}, {4})
        .addEdge(2, {2, 0}, {1, 0}, {4})
        .addEdge(4, {1, 0}, {1, 1}, {5})
        .addEdge(5, {1, 1}, {2, 1}, {6})
        .addEdge(6, {2, 1}, {2, 0}, {2})
        .addEdge(8, {4, 0}, {3, 1}, {9})
        .addEdge(9, {3, 1}, {4, 1}, {10})
        .addEdge(10, {4, 1}, {4, 0}, {8});
    RoadNetworkData roadNetwork = graphCreator.createGraph();

    auto components = getFilteredConnectedComponentsOfTargets(
        roadNetwork, {2, 4, 5, 6, 8, 9, 10});
    EXPECT_EQ(components.size(), 1u);
    std::set<EdgeId> result(components[0].begin(), components[0].end());
    EXPECT_THAT(result, testing::UnorderedElementsAre(2, 4, 5, 6));

}

TEST(connected_components_tests, test_barrier_area)
{
    //  graph:
    // --5->       --9-->
    // ^    |      ^   /
    // |    6      |  10
    // 4    |      8 /
    // |    v      |v
    // <-2-- --3->

    // Components are connected, but not strongly
    // Take only the largest component
    // We consider that all the small components
    // are closed areas

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {1, 0}, {4})
        .addEdge(2, {2, 0}, {1, 0}, {4})
        .addEdge(3, {2, 0}, {3, 0}, {8})
        .addEdge(4, {1, 0}, {1, 1}, {5})
        .addEdge(5, {1, 1}, {2, 1}, {6})
        .addEdge(6, {2, 1}, {2, 0}, {2, 3})
        .addEdge(8, {3, 0}, {3, 1}, {9})
        .addEdge(9, {3, 1}, {4, 1}, {10})
        .addEdge(10, {4, 1}, {3, 0}, {8});
    RoadNetworkData roadNetwork = graphCreator.createGraph();

    auto components = getFilteredConnectedComponentsOfTargets(
        roadNetwork, {2, 3, 4, 5, 6, 8, 9, 10});
    EXPECT_EQ(components.size(), 1u);
    std::set<EdgeId> result(components[0].begin(), components[0].end());
    EXPECT_THAT(result, testing::UnorderedElementsAre(2, 4, 5, 6));
}

TEST(connected_components_tests, test_barrier_area2)
{
    //  graph:
    // ----------------------
    // |                    |
    // | <-6--<-5--<--4--   |
    // | --1->--2->|--3->   |
    // |                    |
    // |  polygon           |
    // ----------------------

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {1, 0}, {2})
        .addEdge(2, {1, 0}, {2, 0}, {5})
        .addEdge(5, {2, 0}, {1, 0}, {6})
        .addEdge(6, {1, 0}, {0, 0}, {1})
        .addEdge(3, {2, 0}, {3, 0}, {4})
        .addEdge(4, {3, 0}, {2, 0}, {3, 5});
    RoadNetworkData roadNetwork = graphCreator.createGraph();
    auto strongConComponents = getStronglyConnectedComponents(roadNetwork);
    EXPECT_EQ(strongConComponents[1], strongConComponents[2]);
    EXPECT_EQ(strongConComponents[1], strongConComponents[5]);
    EXPECT_EQ(strongConComponents[1], strongConComponents[6]);

    EXPECT_NE(strongConComponents[3], strongConComponents[1]);
    EXPECT_EQ(strongConComponents[4], strongConComponents[3]);

    auto components = getFilteredConnectedComponentsOfTargets(
        roadNetwork, {1, 2, 3, 4, 5, 6});
    EXPECT_EQ(components.size(), 1u);
    std::set<EdgeId> result(components[0].begin(), components[0].end());
    EXPECT_THAT(result, testing::UnorderedElementsAre(1, 2, 5, 6));
}

TEST(connected_components_tests, test_two_connected_components)
{
    //  graph:
    // --5-> <-7-- --9-->
    // ^    |      ^   /
    // |    6      |  10
    // 4    |      8 /
    // |    v      |v
    // <-2-- --3->

    // edges 3 and 7 are not targets

    // There are two separate strongly connected components
    // of target
    // We should generate tasks for each component

    GraphCreator graphCreator;
    graphCreator
        .addEdge(2, {2, 0}, {1, 0}, {4})
        .addEdge(3, {2, 0}, {3, 0}, {8}, false)
        .addEdge(4, {1, 0}, {1, 1}, {5})
        .addEdge(5, {1, 1}, {2, 1}, {6})
        .addEdge(6, {2, 1}, {2, 0}, {2, 3})
        .addEdge(7, {3, 1}, {2, 1}, {6}, false)
        .addEdge(8, {3, 0}, {3, 1}, {7, 9})
        .addEdge(9, {3, 1}, {4, 1}, {10})
        .addEdge(10, {4, 1}, {3, 0}, {8});
    RoadNetworkData roadNetwork = graphCreator.createGraph();

    auto weaklyConComponentsOfTargets = getWeaklyConnectedComponents(
        roadNetwork,
        [&] (EdgeId edgeId) {return roadNetwork.edge(edgeId).isTarget;});

    EXPECT_EQ(weaklyConComponentsOfTargets.size(), 7u);

    EXPECT_EQ(weaklyConComponentsOfTargets[2], weaklyConComponentsOfTargets[4]);
    EXPECT_EQ(weaklyConComponentsOfTargets[2], weaklyConComponentsOfTargets[5]);
    EXPECT_EQ(weaklyConComponentsOfTargets[2], weaklyConComponentsOfTargets[6]);

    EXPECT_NE(weaklyConComponentsOfTargets[8], weaklyConComponentsOfTargets[2]);
    EXPECT_EQ(weaklyConComponentsOfTargets[8], weaklyConComponentsOfTargets[9]);
    EXPECT_EQ(weaklyConComponentsOfTargets[8], weaklyConComponentsOfTargets[10]);

    auto components = getFilteredConnectedComponentsOfTargets(
        roadNetwork, {2, 4, 5, 6, 8, 9, 10});
    EXPECT_EQ(components.size(), 2u);
    std::set<EdgeId> result1(components[0].begin(), components[0].end());
    std::set<EdgeId> result2(components[1].begin(), components[1].end());
    if (result1 == std::set<EdgeId>{2, 4, 5, 6}) {
        EXPECT_THAT(result2, testing::UnorderedElementsAre(8, 9, 10));
    } else {
        EXPECT_THAT(result1, testing::UnorderedElementsAre(8, 9, 10));
        EXPECT_THAT(result2, testing::UnorderedElementsAre(2, 4, 5, 6));
    }
}

TEST(connected_components_tests, test_one_big_component)
{
    //  graph:
    // --5-> <-11- --9-->
    // ^    |      ^    |
    // |    6      |    10
    // 4    |      8    |
    // |    v      |    v
    // <-2-- --3-> <-7--

    GraphCreator graphCreator;
    graphCreator
        .addEdge(1, {0, 0}, {1, 0}, {4})
        .addEdge(2, {2, 0}, {1, 0}, {4})
        .addEdge(3, {2, 0}, {3, 0}, {8})
        .addEdge(4, {1, 0}, {1, 1}, {5})
        .addEdge(5, {1, 1}, {2, 1}, {6})
        .addEdge(6, {2, 1}, {2, 0}, {2, 3})
        .addEdge(7, {4, 0}, {3, 0}, {8})
        .addEdge(8, {3, 0}, {3, 1}, {9, 11})
        .addEdge(9, {3, 1}, {4, 1}, {10})
        .addEdge(10, {4, 1}, {4, 0}, {7})
        .addEdge(11, {3, 1}, {2, 1}, {6});
    RoadNetworkData roadNetwork = graphCreator.createGraph();

    auto components = getFilteredConnectedComponentsOfTargets(
        roadNetwork, {2, 3, 4, 5, 6, 7, 8, 9, 10});
    EXPECT_EQ(components.size(), 1u);
    std::set<EdgeId> result(components[0].begin(), components[0].end());
    EXPECT_THAT(result, testing::UnorderedElementsAre(2, 3, 4, 5, 6, 7, 8, 9, 10));
}

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