#include <memory>
#include <print>

#include "wasm/result.hpp"
#include "algorithms/tebd.hpp"
#include "algorithms/dmrg.hpp"
#include "util.hpp"
#include "type.hpp"
#include "wasm/emscripten.hpp"
#include "wasm/range.hpp"


// template<typename Arg>
// void printArgs(Arg&& val) {
//     std::cout << val << std::endl;
// }
// template<typename Arg, typename... Args>
// void printArgs(Arg&& val, Args&&... values) {
//     std::cout << val << ", ";
//     printArgs(std::forward<Args>(values)...);
// }

// template<typename... Args>
// void processDebug(Args&&... args) {
//     std::tuple<typename transform_type<Args>::type...> results = std::make_tuple(processType(std::forward<Args>(args))...);
//     std::apply([](auto&&... args){ printArgs(args...); }, results);
// }

/**
 * @brief Return a string that was set at compilation using the BUILD_OPTIONS_STRING variable
 */
std::string getBuildOptions();


enum class TEBD_TYPE {
    BRICKWALL = 0,
    SECOND_ORDER = 1,
};

inline bool post_update_sleep(const MPS<dtype>&, unsigned) {
    emscripten_sleep(0);
    return false;
};
static postUpdateFunction_t<dtype> post_update_sleep_F(post_update_sleep);

// TEBD
template<typename Model>
void runTEBD(MPS<dtype>& psi, const std::shared_ptr<Model> model, unsigned nSteps, double dtReal, double dtImag, double _eps, unsigned chiMax, TEBD_TYPE type, Results<Model>& result) {
    dtype eps(_eps);
    dtype dt(dtReal, dtImag);
    // 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");
    }
    std::println("Running TEBD type={}, nSteps={}, dt={}, eps={}, chiMax={}", static_cast<int>(type), nSteps, dt, eps, chiMax);

    auto F = result.getPostUpdateFunction(model);
    // initial values now
    F(psi, 0);

    const Bonds<dtype> bonds = tebd::getTimeEvolutionMatrices(model->getH_BondsMatrices(), dt, model->localDim);
    if (type == TEBD_TYPE::BRICKWALL) {
        tebd::runBrickwall(psi, bonds, chiMax, eps, nSteps, F);
    } else if (type == TEBD_TYPE::SECOND_ORDER) {
        const Bonds<dtype> bondsHalf = tebd::getTimeEvolutionMatrices(model->getH_BondsMatrices(), dt/2., model->localDim);
        tebd::runSecondOrder(psi, bonds, bondsHalf, chiMax, eps, nSteps, F);
    } else {
        throw std::invalid_argument(std::format("Got invalid TEBD_TYPE: {}", static_cast<int>(type)));
    }
    std::println("Chis: {}", psi.getChi());
    std::println("Collapse: {}", psi.collapse(0));
    result.finalize(model, psi);
};

/// TODO everything
template<typename Model, typename... ModelArgs>
void runTEBD_evolve(const MPS<dtype>& psiI, unsigned nSteps, double dtReal, double dtImag, double _eps, unsigned chiMax, TEBD_TYPE type, unsigned nValues, Results<Model>& result, ModelArgs&&... args) {
    dtype eps(_eps);
    dtype dt(dtReal, dtImag);
    // sanity checks
    throwIfNaN(eps, "eps");
    throwIfNaN(dt, "dt");
    throwIfLessThan(nSteps, 2u, "nSteps");
    throwIfLessThan(chiMax, 1u, "chiMax");
    std::println("Running TEBD type={}, nSteps={}, dt={}, eps={}, chiMax={}", static_cast<int>(type), nSteps, dt, eps, chiMax);
    MPS<dtype> psi(psiI);

    for (unsigned i = 0; i < nValues; i++) {
        std::shared_ptr<Model> model = makeModel<Model>(std::forward<ModelArgs>(args)...);
        std::println("{:02} - Model: {}", i, *model);
        // std::println("New psi: value at B[2](0, 1, 0)={}", psi.B(2)(0, 1, 0));
        // printTensor(psi.B(0), "psi.B[0]");
        // printTensor(psi.B(1), "psi.B[1]");
        const Bonds<dtype> bonds = tebd::getTimeEvolutionMatrices(model->getH_BondsMatrices(), dt, model->localDim);
        tebd::runBrickwall(psi, bonds, chiMax, eps, 1, post_update_sleep_F);
        result.postUpdateFunction(model, psi);
    }
}


template<typename Model, typename... ModelArgs>
void runTEBD_phaseDiagram(const MPS<dtype>& psiI, unsigned nSteps, double dtReal, double dtImag, double _eps, unsigned chiMax, TEBD_TYPE type, unsigned nValues, Results<Model>& result, ModelArgs&&... args) {
    dtype eps(_eps);
    dtype dt(dtReal, dtImag);
    // sanity checks
    throwIfNaN(eps, "eps");
    throwIfNaN(dt, "dt");
    throwIfLessThan(nSteps, 2u, "nSteps");
    throwIfLessThan(chiMax, 1u, "chiMax");
    std::println("Running TEBD type={}, nSteps={}, dt={}, eps={}, chiMax={}", static_cast<int>(type), nSteps, dt, eps, chiMax);

    for (unsigned i = 0; i < nValues; i++) {
        std::shared_ptr<Model> model = makeModel<Model>(std::forward<ModelArgs>(args)...);
        std::println("{:02} - Model: {}", i, *model);
        MPS<dtype> psi(psiI);
        // std::println("New psi: value at B[2](0, 1, 0)={}", psi.B(2)(0, 1, 0));
        // printTensor(psi.B(0), "psi.B[0]");
        // printTensor(psi.B(1), "psi.B[1]");
        const Bonds<dtype> bonds = tebd::getTimeEvolutionMatrices(model->getH_BondsMatrices(), dt, model->localDim);
        tebd::runBrickwall(psi, bonds, chiMax, eps, nSteps, post_update_sleep_F);
        result.postUpdateFunction(model, psi);
    }
}


// DMRG
template<typename Model>
void runDMRG(MPS<dtype>& psi, const std::shared_ptr<Model> model, unsigned nSweeps, double _eps, unsigned chiMax, Results<Model>& result) {
    dtype eps(_eps);
    // sanity checks
    throwIfNaN(eps, "eps");
    throwIfLessThan(nSweeps, 1u, "nSweeps");
    throwIfLessThan(chiMax, 1u, "chiMax");
    if (psi.getL() != model->L) {
        throw std::runtime_error("psi and model have different L");
    }
    std::println("Running DMRG nSweeps={}, eps={}, chiMax={}", nSweeps, eps, chiMax);

    auto F = result.getPostUpdateFunction(model);
    // initial values now
    F(psi, 0);

    dmrg::run(psi, model->getH_MPO(), chiMax, eps, nSweeps, F, post_update_sleep_F);
    std::println("Chis: {}", psi.getChi());
    std::println("Collapse: {}", psi.collapse(0));
    result.finalize(model, psi);
};


template<typename Model, typename... ModelArgs>
void runDMRG_phaseDiagram(const MPS<dtype>& psiI, unsigned nSweeps, double _eps, unsigned chiMax, unsigned nValues, Results<Model>& result, ModelArgs&&... args) {
    dtype eps(_eps);
    // sanity checks
    throwIfNaN(eps, "eps");
    throwIfLessThan(nSweeps, 1u, "nSweeps");
    throwIfLessThan(chiMax, 1u, "chiMax");
    std::println("Running DMRG nSweeps={}, eps={}, chiMax={}", nSweeps, eps, chiMax);

    for (unsigned i = 0; i < nValues; i++) {
        std::shared_ptr<Model> model = makeModel<Model>(std::forward<ModelArgs>(args)...);
        std::println("{:02} - Model: {}", i, *model);
        MPS<dtype> psi(psiI);
        dmrg::run(psi, model->getH_MPO(), chiMax, eps, nSweeps, post_update_sleep_F, post_update_sleep_F);
        result.postUpdateFunction(model, psi);
        std::println("Collapse: {}", psi.collapse(0));
    }
}
