#include "graph.hpp"
#include <algorithm>
#include <iostream>
#include <random>
#include <vector>

//
// Assumes degree sequence does NOT contain any zeroes!
//
// method2 = true  -> take highest degree and finish its pairing completely
// method2 = false -> take new highest degree after every pairing
template <typename RNG>
bool constrainedConfigurationModel(DegreeSequence &ds, Graph &g, RNG &rng,
                                   bool method2) {
    // Similar to Havel-Hakimi but instead of pairing up with the highest ones
    // that remain, simply pair up with random ones
    unsigned int n = ds.size();

    // degree, vertex index
    std::vector<std::pair<unsigned int, unsigned int>> degrees(n);
    for (unsigned int i = 0; i < n; ++i) {
        degrees[i].first = ds[i];
        degrees[i].second = i;
    }

    // This should make the random pairing faster
    // More likely to pair up with something at the end of an array
    // So erasing that element from the array then requires less moves
    std::sort(
            degrees.begin(), degrees.end(),
            [](const auto &p1, const auto &p2) { return p1.first < p2.first; });

    // remaining half-edges , iterator into `degrees`
    std::vector<std::pair<unsigned int, decltype(degrees.begin())>> available;
    available.reserve(n);

    // Clear the graph
    g.reset(n);

    while (!degrees.empty()) {
        // Get the highest degree:
        // If there are multiple highest ones, we pick the first one
        unsigned int dmax = 0;
        auto uIter = degrees.begin();
        for (auto iter = degrees.begin(); iter != degrees.end(); ++iter) {
            if (iter->first >= dmax) {
                dmax = iter->first;
                uIter = iter;
            }
        }

        if (dmax > degrees.size() - 1)
            return false;

        if (dmax == 0) {
            std::cerr << "ERROR 1 in CCM.\n";
            return false;
        }

        unsigned int u = uIter->second;

        // Pair with a random vertex that is not u itself and to which
        // u has not paired yet!
        // Uniform over all available half-edges, not available vertices
        unsigned int availableEdges = 0;
        available.clear();
        for (auto iter = degrees.begin(); iter != degrees.end(); ++iter) {
            if (iter->second != u && !g.hasEdge({u, iter->second})) {
                availableEdges += iter->first;
                available.push_back(std::make_pair(iter->first, iter));
            }
        }
        if (availableEdges == 0)
            return false;

        if (method2) {
            // Now assign randomly to the remaining vertices
            // weighted by the number of half-edges they have left
            while (dmax--) {
                std::uniform_int_distribution<> distr(1, availableEdges);
                // Go from edge index to number. We should really have a proper
                // data structure like some cumulative tree
                unsigned int halfEdge = distr(rng);
                unsigned int cumulative = 0;
                auto vIter = uIter;
                // Reverse has higher probability to be fast because high
                // degrees are near the end
                for (auto iter = available.rbegin(); iter != available.rend();
                     ++iter) {
                    cumulative += iter->first;
                    if (halfEdge <= cumulative) {
                        vIter = iter->second;
                        availableEdges -= iter->first;
                        // Reverse iterators are tricky
                        available.erase(std::next(iter).base());
                        break;
                    }
                }
                if (vIter == uIter)
                    std::cerr << "ERROR 1 in CCM2.\n";
                if (vIter->first == 0)
                    std::cerr << "ERROR 2 in CCM2.\n";
                if (!g.addEdge({u, vIter->second}))
                    std::cerr << "ERROR 3 in CCM2.\n";
                uIter->first--;
                vIter->first--;
            }
            for (auto iter = degrees.begin(); iter != degrees.end();)
                if (iter->first == 0)
                    iter = degrees.erase(iter);
                else
                    ++iter;
        } else {
            std::uniform_int_distribution<> distr(1, availableEdges);
            // Go from edge index to number. We should really have a proper
            // data structure like some cumulative tree
            auto vIter = uIter;
            unsigned int halfEdge = distr(rng);
            unsigned int cumulative = 0;
            // Reverse has higher probability to be fast because high degrees
            // are near the end
            for (auto iter = available.rbegin(); iter != available.rend();
                 ++iter) {
                cumulative += iter->first;
                if (halfEdge <= cumulative) {
                    vIter = iter->second;
                    break;
                }
            }
            if (vIter == uIter)
                std::cerr << "ERROR 4 in CCM2.\n";
            // pair u to v
            if (vIter->first == 0)
                std::cerr << "ERROR 5 in CCM1.\n";
            if (!g.addEdge({u, vIter->second}))
                std::cerr << "ERROR 6. Could not add edge in CCM1.\n";
            // Purge anything with degree zero
            // Be careful with invalidating the other iterator!
            // Degree of u is always greater or equal to the degree of v
            if (dmax == 1) {
                // Remove both
                // Erasure invalidates all iterators AFTER the erased one
                if (vIter > uIter) {
                    degrees.erase(vIter);
                    degrees.erase(uIter);
                } else {
                    degrees.erase(uIter);
                    degrees.erase(vIter);
                }
            } else {
                // Remove only v if it reaches zero
                uIter->first--;
                vIter->first--;
                if (vIter->first == 0)
                    degrees.erase(vIter);
            }
            // degrees.erase(std::remove_if(degrees.begin(), degrees.end(),
            //                             [](auto x) { return x.first == 0;
            //                             }));
        }
    }
    return true;
}