#include "benchmark.hpp"

#include <print>
#include <vector>
#include <fstream>
#include "mps.hpp"
#include "tebd.hpp"
#include "transverse_field_ising.hpp"
#include "util.hpp"
#include "type.hpp"
#include "stopwatch.hpp"

struct Results {
    std::vector<double> E;
    std::vector<double> M;
    std::vector<double> S;
};

Results _runTEBD(unsigned L, double J, double g, unsigned nSteps, dtype dt, dtype eps, unsigned chiMax) {
    auto psi = initMPS_spinUp<dtype>(L);
    auto model = TransverseFieldIsingModel1D<dtype>(L, J, g);
    // Sanity checks
    throwIfNaN(eps, "eps");
    throwIfNaN(dt, "dt");
    throwIfLessThan(nSteps, 1u, "nSteps");
    throwIfLessThan(chiMax, 1u, "chiMax");
    if (psi.getL() != model.L) {
        throw std::runtime_error("psi and model have different L");
    }
    Results result; 
    result.E.reserve(nSteps+1);
    result.M.reserve(nSteps+1);
    result.S.reserve(nSteps+1);

    e::Tensor<dtype, 2> sigma_z(2, 2);
    sigma_z.setValues({{1, 0}, {0, -1}});

    size_t middleBond = model.L / 2;
    // collect M over time in an array
    auto F = postUpdateFunction_t<dtype>([&result, &sigma_z, &model, middleBond](const MPS<dtype>& psi, unsigned i){
        result.M.push_back(psi.getSiteExpectationValueSum(sigma_z).real());
        result.E.push_back(model.energy(psi).real());
        result.S.push_back(psi.getEntanglementEntropy(middleBond).real());
        return false;
    });
    // initial values now
    F(psi, 0);

    const Bonds<dtype> bonds = tebd::getTimeEvolutionMatrices(model.getH_BondsMatrices(), dt, model.localDim);
    tebd::runBrickwall(psi, bonds, chiMax, eps, nSteps, F);

    return result;
};

template<typename T, typename... Args>
double _runBenchmark(std::function<T(Args...)> f, Args... args) {
    Stopwatch watch;
    watch.start();
    T result = f(std::forward<Args>(args)...);;
    double time = watch.stop();
    return time;
}

int runBenchmark_oneAtATime() {
    // testing parameters
    std::vector<unsigned> chiMaxs{5, 10, 15, 25, 30};
    std::vector<dtype> epss{1e-5, 1e-15, 1e-20};
    std::vector<unsigned> Ls{4, 8, 16, 20, 24, 28, 32};
    std::vector<unsigned> nStepss{400, 600, 800, 1000};

    std::map<std::string, double> times;

    // default parameters
    unsigned L = 12;
    double J = 1;
    double g = 0.7;
    unsigned nSteps = 200;
    dtype dt = 0.05;
    dtype eps = 1e-10;
    unsigned chiMax = 20;

    // get a unique string including all test parameters
    std::function<Results (unsigned, double, double, unsigned, dtype, dtype, unsigned)> f(&_runTEBD);
    auto getName = [](unsigned L, unsigned nSteps, dtype eps, unsigned chiMax) {
        return std::format("L={:02}_nSteps={:04}_eps={:e}_chiMax={:02}", L, nSteps, eps.real(), chiMax);
    };
    // test default
    {
        std::string name = getName(L, nSteps, eps, chiMax);
        std::println("Running {}", name);
        double t = _runBenchmark(f, L, J, g, nSteps, dt, eps, chiMax);
        times[name] = t;
    }

    for (auto _chiMax : chiMaxs) {
        std::string name = getName(L, nSteps, eps, _chiMax);
        std::println("Running {}", name);
        double t = _runBenchmark(f, L, J, g, nSteps, dt, eps, _chiMax);
        times[name] = t;
    }
    for (auto _eps : epss) {
        std::string name = getName(L, nSteps, _eps, chiMax);
        std::println("Running {}", name);
        double t = _runBenchmark(f, L, J, g, nSteps, dt, _eps, chiMax);
        times[name] = t;
    }
    for (auto _L : Ls) {
        std::string name = getName(_L, nSteps, eps, chiMax);
        std::println("Running {}", name);
        double t = _runBenchmark(f, _L, J, g, nSteps, dt, eps, chiMax);
        times[name] = t;
    }
    for (auto _nSteps : nStepss) {
        std::string name = getName(L, _nSteps, eps, chiMax);
        std::println("Running {}", name);
        double t = _runBenchmark(f, L, J, g, _nSteps, dt, eps, chiMax);
        times[name] = t;
    }
    
    std::ofstream file("benchmark.csv", std::ios::app);
    if (!file.is_open()) {
        std::cerr << "Error opening file.\n";
        return 1;
    }
    for (auto [name, t] : times) {
        std::println("{}: {:02f}", name, t);
        file << std::format("\"{}\",{:04f}\n", name, t);
    }
    file.close();

    return 0;
}


int runBenchmark_all() {
    // testing parameters
    // 6 * 4 * 7 * 5 = 840
    // ~10s each => 9000s => 2h30m
    std::vector<unsigned> chiMaxs{5, 15, 25, 35};
    std::vector<dtype> epss{1e-5, 1e-10, 1e-15, 1e-20};
    // std::vector<unsigned> Ls{4, 8, 12, 16, 20, 24, 28, 32};
    std::vector<unsigned> Ls{32, 28, 24, 20, 16, 12, 8, 4};

    // std::vector<unsigned> nStepss{200, 400, 600, 800, 1000};
    std::vector<unsigned> nStepss{200, 400, 600};

    std::map<std::string, double> times;

    // default parameters
    double J = 1;
    double g = 0.7;
    dtype dt = 0.05;

    // get a unique string including all test parameters
    std::function<Results (unsigned, double, double, unsigned, dtype, dtype, unsigned)> f(&_runTEBD);
    auto getName = [](unsigned L, unsigned nSteps, dtype eps, unsigned chiMax) {
        return std::format("L={:02}_nSteps={:04}_eps={:e}_chiMax={:02}", L, nSteps, eps.real(), chiMax);
    };

    for (auto L : Ls) {
        for (auto nSteps : nStepss) {
            for (auto eps : epss) {
                for (auto chiMax : chiMaxs) {
                    std::string name = getName(L, nSteps, eps, chiMax);
                    std::println("Running {}", name);
                    double t = _runBenchmark(f, L, J, g, nSteps, dt, eps, chiMax);
                    times[name] = t;
                }
            }
        }
    }

    std::ofstream file("benchmark.csv", std::ios::app);
    if (!file.is_open()) {
        std::cerr << "Error opening file.\n";
        return 1;
    }
    for (auto [name, t] : times) {
        std::println("{}: {:02f}", name, t);
        file << std::format("\"{}\",{:04f}\n", name, t);
    }
    file.close();

    return 0;
}

#ifdef EMSCRIPTEN
#include <emscripten/bind.h>
namespace ems = emscripten;
EMSCRIPTEN_BINDINGS(benchmark) {
    ems::function("runBenchmark_oneAtATime", &runBenchmark_oneAtATime);
    ems::function("runBenchmark_all", &runBenchmark_all);
};
#endif