/*
 * Unit tests for data structures.
 *
 * Depends on the Boost.Test unit test framework.
 * Tested with Boost v1.74, available from https://www.boost.org/
 */

#include <algorithm>
#include <climits>
#include <map>
#include <memory>
#include <random>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

#define BOOST_TEST_MODULE datastruct
#include <boost/test/unit_test.hpp>

#include "datastruct.h"


/* **********  Test DisjointSetNode  ********** */

BOOST_AUTO_TEST_SUITE(test_disjoint_set)

BOOST_AUTO_TEST_CASE(test_single)
{
    using Node = DisjointSetNode<int>;
    Node a(1);
    BOOST_TEST(a.find() == 1);
    a.set_label(2);
    BOOST_TEST(a.find() == 2);
}

BOOST_AUTO_TEST_CASE(test_simple)
{
    using Node = DisjointSetNode<int>;
    Node a(1);
    Node b(2);
    Node c(3);
    Node* m = a.merge(&b);
    m->set_label(10);
    BOOST_TEST(a.find() == 10);
    BOOST_TEST(b.find() == 10);
    BOOST_TEST(c.find() == 3);
    m = m->merge(&c);
    m->set_label(11);
    BOOST_TEST(a.find() == 11);
    BOOST_TEST(c.find() == 11);
    a.detach(1);
    b.detach(2);
    c.detach(3);
    BOOST_TEST(a.find() == 1);
    BOOST_TEST(b.find() == 2);
    BOOST_TEST(c.find() == 3);
}

BOOST_AUTO_TEST_CASE(test_multilevel)
{
    using Node = DisjointSetNode<int>;

    std::unique_ptr<Node> nodes[27];
    for (int i = 0; i < 27; ++i) {
        nodes[i].reset(new Node(i));
    }

    std::vector<Node*> level1;
    for (int i = 0; i < 9; ++i) {
        Node* m = nodes[3*i].get();
        for (int k = 1; k < 3; ++k) {
            m = m->merge(nodes[3*i+k].get());
        }
        m->set_label(100 + i);
        level1.push_back(m);
    }

    std::vector<Node*> level2;
    for (int i = 0; i < 3; ++i) {
        Node* m = level1[3*i];
        for (int k = 1; k < 3; ++k) {
            m = m->merge(level1[3*i+k]);
        }
        m->set_label(200 + i);
        level2.push_back(m);
    }

    Node* m = level2[0];
    for (int k = 1; k < 3; ++k) {
        m = m->merge(level2[k]);
    }
    m->set_label(300);

    for (int i = 0; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == 300);
    }

    for (int i = 0; i < 3; ++i) {
        level2[i]->detach(200 + i);
    }

    for (int i = 0; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == 200 + i / 9);
    }

    for (int i = 0; i < 6; ++i) {
        level1[i]->detach(100 + i);
    }

    for (int i = 0; i < 18; ++i) {
        BOOST_TEST(nodes[i]->find() == 100 + i / 3);
    }
    for (int i = 18; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == 202);
    }

    for (int i = 6; i < 9; ++i) {
        level1[i]->detach(100 + i);
    }

    for (int i = 0; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == 100 + i / 3);
    }

    for (int i = 6; i < 27; ++i) {
        nodes[i]->detach(i);
    }

    for (int i = 0; i < 6; ++i) {
        BOOST_TEST(nodes[i]->find() == 100 + i / 3);
    }
    for (int i = 6; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == i);
    }

    for (int i = 0; i < 6; i++) {
        nodes[i]->detach(i);
    }

    for (int i = 0; i < 27; ++i) {
        BOOST_TEST(nodes[i]->find() == i);
    }
}

BOOST_AUTO_TEST_CASE(test_random)
{
    using Node = DisjointSetNode<int>;

    std::mt19937 rng(12345);

    const int num_nodes = 1000;
    std::unique_ptr<Node> nodes[num_nodes];
    for (int i = 0; i < num_nodes; ++i) {
        nodes[i].reset(new Node(i));
    }

    std::unordered_map<int, Node*> blossoms;
    std::unordered_map<int, std::vector<int>> sub_blossoms;
    std::vector<int> top_blossoms;

    for (int i = 0; i < num_nodes; ++i) {
        blossoms[i] = nodes[i].get();
        top_blossoms.push_back(i);
    }

    int next_blossom = num_nodes;

    auto make_blossom = [&]() {
        int b = next_blossom;
        ++next_blossom;

        int nsub = 1 + 2 * std::uniform_int_distribution<>(1, 4)(rng);
        std::vector<int> subs(nsub);

        for (int i = 0; i < nsub; ++i) {
            int p = std::uniform_int_distribution<>(0, top_blossoms.size() - 1)(rng);
            subs[i] = top_blossoms[p];
            top_blossoms.erase(top_blossoms.begin() + p);
        }

        Node *m = blossoms[subs[0]];
        for (int i = 1; i < nsub; ++i) {
            m = m->merge(blossoms[subs[i]]);
        }
        m->set_label(b);

        blossoms[b] = m;
        sub_blossoms[b] = std::move(subs);
        top_blossoms.push_back(b);
    };

    auto expand_blossom = [&](int b) {
        top_blossoms.erase(
            std::find(top_blossoms.begin(), top_blossoms.end(), b));
        for (int t : sub_blossoms[b]) {
            blossoms[t]->detach(t);
            top_blossoms.push_back(t);
        }
        blossoms.erase(b);
        sub_blossoms.erase(b);
    };

    auto check_membership = [&](int b, int label) {
        std::vector<int> q;
        q.push_back(b);
        while (! q.empty()) {
            b = q.back();
            q.pop_back();
            if (b < num_nodes) {
                BOOST_TEST(nodes[b]->find() == label);
            } else {
                for (int t : sub_blossoms[b]) {
                    q.push_back(t);
                }
            }
        }
    };

    for (int k = 0; k < 100; ++k) {
        make_blossom();
    }

    for (int b : top_blossoms) {
        check_membership(b, b);
    }

    std::vector<int> top_groups;
    for (int b : top_blossoms) {
        if (b >= num_nodes) {
            top_groups.push_back(b);
        }
    }
    std::shuffle(top_groups.begin(), top_groups.end(), rng);
    for (int k = 0; k < 50; ++k) {
        expand_blossom(top_groups[k]);
    }
    top_groups.clear();

    for (int b : top_blossoms) {
        check_membership(b, b);
    }

    for (int k = 0; k < 50; ++k) {
        make_blossom();
    }

    for (int b : top_blossoms) {
        check_membership(b, b);
    }

    for (int b : top_blossoms) {
        if (b >= num_nodes) {
            top_groups.push_back(b);
        }
    }
    std::shuffle(top_groups.begin(), top_groups.end(), rng);
    for (int b : top_groups) {
        expand_blossom(b);
    }

    for (int b : top_blossoms) {
        check_membership(b, b);
    }
}

BOOST_AUTO_TEST_SUITE_END()


/* **********  Test SplittablePriorityQueue  ********** */

template <typename PrioType, typename DataType>
static void check_min_elem(const std::pair<PrioType, DataType>& pair,
                           PrioType prio,
                           const DataType& data)
{
    BOOST_TEST(pair.first == prio);
    BOOST_TEST(pair.second == data);
}


BOOST_AUTO_TEST_SUITE(test_splittable_queue)

BOOST_AUTO_TEST_CASE(test_single)
{
    using Queue = SplittablePriorityQueue<int, int>;
    Queue q;

    BOOST_TEST(q.empty() == true);

    Queue::Node n;
    q.insert(&n, 3, 4, 101);

    BOOST_TEST(q.empty() == false);
    check_min_elem(q.find_min(), 4, 101);

    q.update(&n, 5, 102);
    check_min_elem(q.find_min(), 4, 101);

    q.update(&n, 3, 103);
    check_min_elem(q.find_min(), 3, 103);

    q.clear();
    BOOST_TEST(q.empty() == true);
}

BOOST_AUTO_TEST_CASE(test_simple)
{
    using Queue = SplittablePriorityQueue<int, std::string>;
    Queue q;
    Queue::Node n1, n2, n3, n4, n5, nx;

    q.insert(&n1, 1, 5, "a");
    q.insert(&n2, 2, 6, "b");
    q.insert(&n3, 3, 7, "c");
    q.insert(&n4, 4, 4, "d");
    q.insert(&n5, 5, 3, "e");
    check_min_elem(q.find_min(), 3, std::string("e"));

    q.update(&n1, 4, "f");
    check_min_elem(q.find_min(), 3, std::string("e"));

    q.update(&n3, 2, "h");
    check_min_elem(q.find_min(), 2, std::string("h"));

    Queue q2 = q.split(3);
    check_min_elem(q.find_min(), 4, std::string("f"));
    check_min_elem(q2.find_min(), 2, std::string("h"));

    q.insert(&nx, 3, 1, "x");
    check_min_elem(q.find_min(), 1, std::string("x"));
    check_min_elem(q2.find_min(), 2, std::string("h"));

    q.clear();
    q2.clear();
}

BOOST_AUTO_TEST_CASE(test_split_empty)
{
    using Queue = SplittablePriorityQueue<int, std::string>;
    Queue q;
    Queue q2 = q.split(10);
    BOOST_TEST(q.empty() == true);
    BOOST_TEST(q2.empty() == true);
}

BOOST_AUTO_TEST_CASE(test_split_oneway)
{
    using Queue = SplittablePriorityQueue<int, std::string>;
    Queue q;
    Queue::Node n4, n5, n6;
    q.insert(&n4, 4, 3, "a");
    q.insert(&n5, 5, 4, "b");
    q.insert(&n6, 6, 2, "c");
    Queue q2 = q.split(7);
    BOOST_TEST(q.empty() == false);
    BOOST_TEST(q2.empty() == true);
    check_min_elem(q.find_min(), 2, std::string("c"));
    q.clear();

    q.insert(&n4, 4, 3, "a");
    q.insert(&n5, 5, 4, "b");
    q.insert(&n6, 6, 2, "c");
    q2 = q.split(4);
    BOOST_TEST(q.empty() == true);
    BOOST_TEST(q2.empty() == false);
    check_min_elem(q2.find_min(), 2, std::string("c"));
    q2.clear();
}

BOOST_AUTO_TEST_CASE(test_larger)
{
    using Queue = SplittablePriorityQueue<int, std::string>;
    Queue q;
    Queue::Node nodes[15];

    q.insert(&nodes[7],   7, 5, "h");
    q.insert(&nodes[6],   6, 4, "g");
    q.insert(&nodes[8],   8, 2, "i");
    q.insert(&nodes[5],   5, 4, "f");
    q.insert(&nodes[9],   9, 6, "j");
    q.insert(&nodes[4],   4, 8, "e");
    q.insert(&nodes[10], 10, 4, "k");
    q.insert(&nodes[3],   3, 5, "d");
    q.insert(&nodes[11], 11, 6, "l");
    q.insert(&nodes[2],   2, 7, "c");
    q.insert(&nodes[12], 12, 8, "m");
    q.insert(&nodes[1],   1, 3, "b");
    q.insert(&nodes[13], 13, 1, "n");
    q.insert(&nodes[0],   0, 9, "a");
    q.insert(&nodes[14], 14, 7, "o");

    check_min_elem(q.find_min(), 1, std::string("n"));

    Queue q2 = q.split(10);
    check_min_elem(q.find_min(), 2, std::string("i"));
    check_min_elem(q2.find_min(), 1, std::string("n"));

    Queue q1 = q.split(5);
    check_min_elem(q.find_min(), 3, std::string("b"));
    check_min_elem(q1.find_min(), 2, std::string("i"));
    check_min_elem(q2.find_min(), 1, std::string("n"));

    q.clear();
    q1.clear();
    q2.clear();
}

BOOST_AUTO_TEST_CASE(test_random)
{
    using Queue = SplittablePriorityQueue<int, int>;

    std::mt19937 rng(23456);
    std::uniform_int_distribution<> index_distribution(0, 1000000);
    std::uniform_int_distribution<> prio_distribution(0, 1000000);

    Queue q;
    std::map<int, std::tuple<std::unique_ptr<Queue::Node>, int, int>> elems;
    int next_data = 1;

    for (int i = 0; i < 200; ++i) {

        // Insert stuff into the queue.
        for (int k = 0; k < 1000; ++k) {
            int idx = index_distribution(rng);
            int prio = prio_distribution(rng);
            int data = next_data;
            ++next_data;
            auto it = elems.find(idx);
            if (it != elems.end()) {
                auto& v = it->second;
                Queue::Node* nptr = std::get<0>(v).get();
                int pprio = std::get<1>(v);
                q.update(nptr, prio, data);
                if (prio < pprio) {
                    std::get<1>(v) = prio;
                    std::get<2>(v) = data;
                }
            } else {
                std::unique_ptr<Queue::Node> nptr(new Queue::Node);
                q.insert(nptr.get(), idx, prio, data);
                elems[idx] = std::make_tuple(std::move(nptr), prio, data);
            }
        }

        // Check min element.
        int min_prio = INT_MAX;
        int min_data = 0;
        for (const auto& v : elems) {
            int prio = std::get<1>(v.second);
            int data = std::get<2>(v.second);
            if (prio < min_prio) {
                min_prio = prio;
                min_data = data;
            }
        }

        check_min_elem(q.find_min(), min_prio, min_data);

        // Split the queue.
        int threshold = index_distribution(rng);
        Queue q2 = q.split(threshold);

        // Keep one queue and discard the other.
        if (rng() % 2 == 0) {
            q.clear();
            q = std::move(q2);
            elems.erase(elems.begin(), elems.lower_bound(threshold));
        } else {
            q2.clear();
            elems.erase(elems.lower_bound(threshold), elems.end());
        }

        // Check min element.
        min_prio = INT_MAX;
        min_data = 0;
        for (const auto& v : elems) {
            int prio = std::get<1>(v.second);
            int data = std::get<2>(v.second);
            if (prio < min_prio) {
                min_prio = prio;
                min_data = data;
            }
        }

        if (min_prio < INT_MAX) {
            check_min_elem(q.find_min(), min_prio, min_data);
        }
    }

    q.clear();
}

BOOST_AUTO_TEST_SUITE_END()


/* **********  Test PriorityQueue  ********** */

BOOST_AUTO_TEST_SUITE(test_priority_queue)

BOOST_AUTO_TEST_CASE(test_empty)
{
    using Queue = PriorityQueue<int, std::string>;
    Queue q;
    BOOST_TEST(q.empty() == true);
}

BOOST_AUTO_TEST_CASE(test_single)
{
    using Queue = PriorityQueue<int, std::string>;
    Queue q;

    Queue::Node n1;
    q.insert(&n1, 5, "a");

    BOOST_TEST(q.empty() == false);
    check_min_elem(q.find_min(), 5, std::string("a"));

    q.update(&n1, 3, "a");
    check_min_elem(q.find_min(), 3, std::string("a"));

    q.remove(&n1);
    BOOST_TEST(q.empty() == true);
}

BOOST_AUTO_TEST_CASE(test_simple)
{
    using Queue = PriorityQueue<int, char>;
    Queue q;
    Queue::Node nodes[10];

    q.insert(&nodes[0], 9, 'a');
    check_min_elem(q.find_min(), 9, 'a');

    q.insert(&nodes[1], 4, 'b');
    check_min_elem(q.find_min(), 4, 'b');

    q.insert(&nodes[2], 7, 'c');
    check_min_elem(q.find_min(), 4, 'b');

    q.insert(&nodes[3], 5, 'd');
    check_min_elem(q.find_min(), 4, 'b');

    q.insert(&nodes[4], 8, 'e');
    check_min_elem(q.find_min(), 4, 'b');

    q.insert(&nodes[5], 6, 'f');
    check_min_elem(q.find_min(), 4, 'b');

    q.insert(&nodes[6], 4, 'g');
    q.insert(&nodes[7], 5, 'h');
    q.insert(&nodes[8], 2, 'i');
    check_min_elem(q.find_min(), 2, 'i');

    q.insert(&nodes[9], 6, 'j');
    check_min_elem(q.find_min(), 2, 'i');

    q.update(&nodes[2], 1, 'c');
    check_min_elem(q.find_min(), 1, 'c');

    q.update(&nodes[4], 3, 'e');
    check_min_elem(q.find_min(), 1, 'c');

    q.remove(&nodes[2]);
    check_min_elem(q.find_min(), 2, 'i');

    q.remove(&nodes[8]);
    check_min_elem(q.find_min(), 3, 'e');

    q.remove(&nodes[4]);
    q.remove(&nodes[1]);
    check_min_elem(q.find_min(), 4, 'g');

    q.remove(&nodes[3]);
    q.remove(&nodes[9]);
    check_min_elem(q.find_min(), 4, 'g');

    q.remove(&nodes[6]);
    check_min_elem(q.find_min(), 5, 'h');

    BOOST_TEST(q.empty() == false);
    q.clear();
    BOOST_TEST(q.empty() == true);
}

BOOST_AUTO_TEST_CASE(test_random)
{
    using Queue = PriorityQueue<int, int>;
    Queue q;

    const int num_elem = 1000;
    std::vector<std::tuple<std::unique_ptr<Queue::Node>, int, int>> elems;
    int next_data = 0;

    std::mt19937 rng(34567);

    auto check = [&q,&elems]() {
        int min_prio, min_data;
        std::tie(min_prio, min_data) = q.find_min();
        int best_prio = INT_MAX;
        bool found = false;
        for (const auto& v : elems) {
            int this_prio = std::get<1>(v);
            int this_data = std::get<2>(v);
            best_prio = std::min(best_prio, this_prio);
            if ((this_prio == min_prio) && (this_data == min_data)) {
                found = true;
            }
        }
        BOOST_TEST(found == true);
        BOOST_TEST(min_prio == best_prio);
    };

    for (int i = 0; i < num_elem; ++i) {
        ++next_data;
        int prio = std::uniform_int_distribution<>(0, 1000000)(rng);
        std::unique_ptr<Queue::Node> nptr(new Queue::Node);
        q.insert(nptr.get(), prio, next_data);
        elems.push_back(std::make_tuple(std::move(nptr), prio, next_data));
        check();
    }

    for (int i = 0; i < 10000; ++i) {
        int p = std::uniform_int_distribution<>(0, num_elem - 1)(rng);
        Queue::Node* node = std::get<0>(elems[p]).get();
        int prio = std::get<1>(elems[p]);
        int data = std::get<2>(elems[p]);
        prio = std::uniform_int_distribution<>(0, prio)(rng);
        q.update(node, prio, data);
        std::get<1>(elems[p]) = prio;
        check();

        p = std::uniform_int_distribution<>(0, num_elem - 1)(rng);
        node = std::get<0>(elems[p]).get();
        q.remove(node);
        elems.erase(elems.begin() + p);
        check();

        ++next_data;
        prio = std::uniform_int_distribution<>(0, 1000000)(rng);
        std::unique_ptr<Queue::Node> nptr(new Queue::Node);
        q.insert(nptr.get(), prio, next_data);
        elems.push_back(std::make_tuple(std::move(nptr), prio, next_data));
        check();
    }

    for (int i = 0; i < num_elem; ++i) {
        int p = std::uniform_int_distribution<>(0, num_elem - 1 - i)(rng);
        Queue::Node* node = std::get<0>(elems[p]).get();
        q.remove(node);
        elems.erase(elems.begin() + p);
        if (! elems.empty()) {
            check();
        }
    }

    BOOST_TEST(q.empty() == true);
}

BOOST_AUTO_TEST_SUITE_END()