#include "exports.hpp"
#include "graph.hpp"
#include "powerlaw.hpp"
#include <algorithm>
#include <fstream>
#include <iostream>
#include <numeric>
#include <random>
#include <vector>

// Its assumed that u,v are distinct.
// Check if all four vertices are distinct
bool edgeConflicts(const Edge& e1, const Edge& e2) {
    return (e1.u == e2.u || e1.u == e2.v || e1.v == e2.u || e1.v == e2.v);
}

class SwitchChain {
  public:
    SwitchChain()
        : mt(std::random_device{}()), permutationDistribution(0.5)
    // permutationDistribution(0, 2)
    {
        // random_device uses hardware entropy if available
        // std::random_device rd;
        // mt.seed(rd());
    }
    ~SwitchChain() {}

    bool initialize(const Graph& gstart) {
        if (gstart.edgeCount() == 0)
            return false;
        g = gstart;
        edgeDistribution.param(
            std::uniform_int_distribution<>::param_type(0, g.edgeCount() - 1));
        return true;
    }

    bool doMove() {
        int e1index, e2index;
        int timeout = 0;
        // Keep regenerating while conflicting edges
        do {
            e1index = edgeDistribution(mt);
            e2index = edgeDistribution(mt);
            if (++timeout % 100 == 0) {
                std::cerr << "Warning: sampled " << timeout
                          << " random edges but they keep conflicting.\n";
            }
        } while (edgeConflicts(g.getEdge(e1index), g.getEdge(e2index)));

        // Consider one of the three possible permutations
        // 1) e1.u - e1.v and e2.u - e2.v (original)
        // 2) e1.u - e2.u and e1.v - e2.v
        // 3) e1.u - e2.v and e1.v - e2.u
        bool switchType = permutationDistribution(mt);
        return g.exchangeEdges(e1index, e2index, switchType);
    }

    Graph g;
    std::mt19937 mt;
    std::uniform_int_distribution<> edgeDistribution;
    //std::uniform_int_distribution<> permutationDistribution;
    std::bernoulli_distribution permutationDistribution;
};

bool greedyConfigurationModel(DegreeSequence& ds, Graph& g, auto& rng) {
    // Same as 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;
    }
    std::vector<decltype(degrees.begin())> available;
    available.reserve(n);

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

    std::uniform_int_distribution<> distr(1, n - 1);

    while (!degrees.empty()) {
        // Get the highest degree:
        unsigned int dmax = 0;
        unsigned int u = 0;
        auto maxIter = degrees.begin();
        for (auto iter = degrees.begin(); iter != degrees.end(); ++iter) {
            if (iter->first >= dmax) {
                dmax = iter->first;
                u = iter->second;
                maxIter = iter;
            }
        }
        // Take the highest degree out of the vector
        degrees.erase(maxIter);

        available.clear();
        for (auto iter = degrees.begin(); iter != degrees.end(); ++iter) {
            if (iter->first)
                available.push_back(iter);
        }

        if (dmax > available.size())
            return false;

        std::shuffle(available.begin(), available.end(), rng);

        // Now assign randomly to the remaining vertices
        // Pick 'cmax' distinct integers between 1 and degrees.size()
        for (unsigned int i = 0; i < dmax; ++i) {
            if (!g.addEdge({u,available[i]->second})) {
                std::cerr << "ERROR. Could not add edge in greedy configuration model.\n";
            }
            available[i]->first--;
        }
    }
    return true;
}

int main() {
    // Generate a random degree sequence
    std::mt19937 rng(std::random_device{}());

    // Goal:
    // Degrees follow a power-law distribution with some parameter tau
    // Expect:  #tri = const * n^{ something }
    // The goal is to find the 'something' by finding the number of triangles
    // for different values of n and tau
    float tauValues[] = {2.1f, 2.2f, 2.3f, 2.4f, 2.5f, 2.6f, 2.7f, 2.8f};

    Graph g;

    std::ofstream outfile("graphdata.m");
    outfile << '{';
    bool outputComma = false;

    for (int numVertices = 200; numVertices <= 1000; numVertices += 100) {
        for (float tau : tauValues) {

            DegreeSequence ds(numVertices);
            powerlaw_distribution degDist(tau, 1, numVertices);
            //std::poisson_distribution<> degDist(12);

            // For a single n,tau take samples over several instances of
            // the degree distribution
            for (int degreeSample = 0; degreeSample < 500; ++degreeSample) {
                // Generate a graph
                // might require multiple tries
                for (int i = 1; ; ++i) {
                    std::generate(ds.begin(), ds.end(),
                                  [&degDist, &rng] { return degDist(rng); });
                    // First make the sum even
                    unsigned int sum = std::accumulate(ds.begin(), ds.end(), 0);
                    if (sum % 2) {
                        continue;
                        // Can we do this: ??
                        ds.back()++;
                    }

                    // Option 1:
                    if (g.createFromDegreeSequence(ds)) {
                        // Test option 2:
                        //int good = 0;
                        //for (int i = 0; i < 100; ++i) {
                        //    if (greedyConfigurationModel(ds, g, rng))
                        //        ++good;
                        //}
                        //std::cerr << "Greedy configuration model success rate: " << good << "%\n";

                        //g.createFromDegreeSequence(ds);
                        break;
                    }

                    // When 10 tries have not worked, output a warning
                    if (i % 10 == 0) {
                        std::cerr << "Warning: could not create graph from "
                                     "degree sequence. Trying again...\n";
                    }
                }

                SwitchChain chain;
                if (!chain.initialize(g)) {
                    std::cerr << "Could not initialize Markov chain.\n";
                    return 1;
                }

                std::cout << "Running n = " << numVertices << ", tau = " << tau
                          << ". \t" << std::flush;

                int mixingTime = (32.0f - 26.0f*(tau - 2.0f)) * numVertices; //40000;
                constexpr int measurements = 50;
                constexpr int measureSkip =
                    200; // Take a sample every ... steps
                int movesDone = 0;

                int triangles[measurements];

                for (int i = 0; i < mixingTime; ++i) {
                    if (chain.doMove())
                        ++movesDone;
                }
                for (int i = 0; i < measurements; ++i) {
                    for (int j = 0; j < measureSkip; ++j)
                        if (chain.doMove())
                            ++movesDone;
                    triangles[i] = chain.g.countTriangles();
                }

                std::cout << movesDone << '/' << mixingTime + measurements * measureSkip
                          << " moves succeeded ("
                          << 100.0f * float(movesDone) /
                                 float(mixingTime + measurements * measureSkip)
                          << "%)." << std::endl;

                if (outputComma)
                    outfile << ',';
                outputComma = true;

                std::sort(ds.begin(), ds.end());
                outfile << '{' << '{' << numVertices << ',' << tau << '}';
                outfile << ',' << triangles << ',' << ds << '}' << std::flush;
            }
        }
    }
    outfile << '}';
    return 0;
}