#include "exports.hpp"
#include "graph.hpp"
#include "graph_powerlaw.hpp"
#include "graph_spectrum.hpp"
#include "switchchain.hpp"
#include <algorithm>
#include <fstream>
#include <iostream>
#include <numeric>
#include <random>
#include <vector>

double getDSTN(const DegreeSequence& ds) {
    std::vector<std::vector<double>> vals(ds.size());
    for (auto& v : vals) {
        v.resize(ds.size(), 0);
    }

    auto D = 0u;
    for (auto d : ds)
        D += d;

    double factor = 1.0 / double(D);

    for (auto i = 0u; i < ds.size(); ++i) {
        for (auto j = i + 1; j < ds.size(); ++j) {
            vals[i][j] = 1.0 - std::exp(-(ds[i] * ds[j] * factor));
        }
    }

    double result = 0.0;
    for (auto i = 0u; i < ds.size(); ++i) {
        for (auto j = i + 1; j < ds.size(); ++j) {
            for (auto k = j + 1; k < ds.size(); ++k) {
                result += vals[i][j] * vals[j][k] * vals[i][k];
            }
        }
    }
    return result;
}

int main(int argc, char* argv[]) {
    // Simulation parameters
    const int numVerticesMin = 1000;
    const int numVerticesMax = 10000;
    const int numVerticesStep = 1000;

    float tauValues[] = {2.1f, 2.2f, 2.3f, 2.4f, 2.5f, 2.6f, 2.7f, 2.8f, 2.9f};

    //const int totalDegreeSamples = 5000;

    auto getMixingTime = [](int n, float tau) {
        return int(50.0f * (50.0f - 30.0f * (tau - 2.0f)) * n);
    };
    constexpr int measurements = 10;
    constexpr int measureSkip = 1000; // Take a sample every ... steps

    // Output file
    std::ofstream outfile;
    if (argc >= 2)
        outfile.open(argv[1]);
    else
        outfile.open("graphdata_canonical_properties.m");
    if (!outfile.is_open()) {
        std::cout << "ERROR: Could not open output file.\n";
        return 1;
    }

    // Output Mathematica-style comment to indicate file contents
    outfile << "(*\n";
    outfile << "n from " << numVerticesMin << " to " << numVerticesMax
            << " step " << numVerticesStep << std::endl;
    outfile << "tauValues: " << tauValues << std::endl;
    outfile << "Canonical degree sequence.\n";
    outfile << "mixingTime: 50 * (50 - 30 (tau - 2)) n\n";
    outfile << "data:\n";
    outfile << "1: {n,tau}\n";
    outfile << "2: avgTriangles\n";
    outfile << "3: edges\n";
    outfile << "4: dstn\n";
    outfile << "5: { HH A, HH L, average A, average L }  where for each there is (average of) {lambda1 , lambda1 - lambda2, lambda1/lambda2}\n";
    outfile << "6: switching successrate after mixing\n";
    outfile << "7: initial HH triangles\n";
    outfile << "*)" << std::endl;

    // Mathematica does not accept normal scientific notation
    outfile << std::fixed;
    outfile << '{';
    bool outputComma = false;

    Graph g;
    for (int numVertices = numVerticesMin; numVertices <= numVerticesMax;
         numVertices += numVerticesStep) {
        for (float tau : tauValues) {
            DegreeSequence ds;
            generateCanonicalPowerlawGraph(numVertices, tau, g, ds);

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

            std::cout << "Running (n,tau) = (" << numVertices << ',' << tau
                      << "). " << std::flush;

            // Mix
            int mixingTime = getMixingTime(numVertices, tau);
            for (int i = 0; i < mixingTime; ++i) {
                chain.doMove();
            }

            std::cout << "Mixing done. " << std::flush;

            std::array<double, 3> HHAspectrum;
            std::array<double, 3> HHLspectrum;
            std::array<double, 3> avgAspectrum;
            std::array<double, 3> avgLspectrum;

            auto getSpectralValues =
                [](const std::vector<float> &s) -> std::array<double, 3> {
                auto l1 = s[s.size() - 1];
                auto l2 = s[s.size() - 2];
                return {l1, l1 - l2, l1 / l2};
            };

            GraphSpectrum gs_start(g);
            GraphSpectrum gs(chain.g);

            HHAspectrum =
                getSpectralValues(gs_start.computeAdjacencySpectrum());
            HHLspectrum =
                getSpectralValues(gs_start.computeLaplacianSpectrum());

            long long trianglesTotal = 0;
            int movesDone = 0;
            avgAspectrum.fill(0);
            avgLspectrum.fill(0);
            for (int i = 0; i < measurements; ++i) {
                for (int j = 0; j < measureSkip; ++j)
                    if (chain.doMove())
                        ++movesDone;
                trianglesTotal += chain.g.countTriangles();
                auto sA = getSpectralValues(gs.computeAdjacencySpectrum());
                auto sL = getSpectralValues(gs.computeLaplacianSpectrum());
                for (auto i = 0u; i < 3; ++i) {
                    avgAspectrum[i] += sA[i];
                    avgLspectrum[i] += sL[i];
                }
            }
            float avgTriangles = float(trianglesTotal) / float(measurements);
            float successrate =
                float(movesDone) / float(measurements * measureSkip);
            for (auto &f : avgAspectrum)
                f /= float(measurements);
            for (auto &f : avgLspectrum)
                f /= float(measurements);

            std::cout << "Measuring done." << std::flush;

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

            outfile << '{' << '{' << numVertices << ',' << tau << '}';
            outfile << ',' << avgTriangles;
            outfile << ',' << g.edgeCount();
            outfile << ',' << getDSTN(ds);
            outfile << ',' << '{' << HHAspectrum;
            outfile << ',' << HHLspectrum;
            outfile << ',' << avgAspectrum;
            outfile << ',' << avgLspectrum;
            outfile << '}';
            outfile << ',' << successrate;
            outfile << ',' << g.countTriangles();
            outfile << '}' << std::flush;

            std::cout << "Output done." << std::endl;
        }
    }
    outfile << '}';
    return 0;
}