#include <print>
#include <iostream>
// #include <array>
#include "bose_hubbard.hpp"
#include "mps.hpp"
#include "transverse_field_ising.hpp"
#include "tebd.hpp"
#include "util.hpp"
#include "type.hpp"
#include "dmrg.hpp"
#include "wasm/run.hpp"

#include "benchmark.hpp"

using namespace std::complex_literals;

void testWasmCode() {
    unsigned L = 6;
    unsigned nMax = 3;
    auto model = std::make_shared<BoseHubbardModel1D<dtype>>(L, dtype(1.), dtype(0.05), dtype(0.5), nMax);
    auto psi = initMPS_nParticles<dtype>(L, nMax, 1);

    double eps = 1e-10;
    unsigned chiMax = 20;
    double dtReal = 0.00;
    double dtImag = 0.05;
    unsigned nSteps = 300;
    ems::val callback;
    BoseHubbardResults result(callback, 0);
    runTEBD(psi, model, nSteps, dtReal, dtImag, eps, chiMax, TEBD_TYPE::BRICKWALL, result);
}

void testSVD() {
    e::Tensor<dtype, 4> t4(2, 3, 4, 5);
    auto data = t4.data();
    for (unsigned i = 0; i < t4.size(); i++) {
        data[i] = i;
    }
    printTensor(t4, "t4");
    auto [A, S, B] = split_truncate_theta<dtype>(t4, 20, 1e-10);
    printTensor(A, "A");
    printTensor(S, "S");
    printTensor(B, "B");
    e::Tensor<dtype, 4> t4reconstructed = A.contract(diagonal(S).contract(B, e::array<iP, 1>{iP{1, 0}}), e::array<iP, 1>{iP{2, 0}});
    // this can only yield the original tensor if the normalization of the schmidt values in the split_truncate_theta function is removed
    printTensor(t4reconstructed, "t4reconstructed");
}


void testBoseHubbardTEBD() {
    unsigned L = 6;
    unsigned nMax = 3;
    auto model = BoseHubbardModel1D<dtype>(L, dtype(1.), dtype(0.05), dtype(0.5), nMax);

    dtype eps = 1e-10;
    unsigned chiMax = 20;
    dtype dt = 0.05i;
    unsigned nSteps = 300;
    std::println("dt={}", dt);
    // std::cout << "leftT:\n" << *(model.getH_Bonds().front()) << std::endl;
    // std::cout << "leftT2:\n" << bondMatrix2Tensor(*(model.getH_BondsMatrices().front()), model.localDim) << std::endl;
    std::cout << "middleT:\n" << *(model.getH_Bonds().at(1)) << std::endl;
    //
    auto bonds = tebd::getTimeEvolutionMatrices(model.getH_BondsMatrices(), dt, model.localDim);
    std::cout << "N:\n" << model.getLadderOperators().getN() << std::endl;
    std::cout << "a:\n" << model.getLadderOperators().getA() << std::endl;
    std::cout << "aDagger:\n" << model.getLadderOperators().getADagger() << std::endl;
    // std::cout << "U_Bond:\n" << *(bonds.front()) << std::endl;
    // std::cout << "H_Bond:\n" << *(bonds.front()) << std::endl;

    // MPS<dtype> psi = initMPS_spinUp<dtype>(L);
    MPS<dtype> psi = initMPS_nParticles<dtype>(L, nMax, 1);

    // postUpdateFunction_t<dtype> fUpdate = postUpdateNoop<dtype>;

    e::TensorMap<const e::Tensor<dtype, 2>> particleNumberOperator(model.getLadderOperators().getN().data(), model.localDim, model.localDim);
    postUpdateFunction_t<dtype> fUpdate = [&model, &particleNumberOperator](const MPS<dtype>& psi, unsigned i) {
        std::println("{:03}: N={:.2f} {}", i, psi.getSiteExpectationValueSum(particleNumberOperator).real() / model.L, psi.getChi());
        return false;
    };
    // auto fUpdate = postUpdateFunction_t<dtype>([](const MPS<dtype>& psi, unsigned i){
    //     std::println("i={:03}: <psi|psi>={}", i, psi.collapse(0));
    // });
    // std::println("E = {}", model.energy(psi).real());
    // std::println("m = {}", psi.getSiteExpectationValues(sigma_z));
    // std::println("chiS: {:b}", psi.getChi());
    // psi.print();
    // std::println("<psi|psi>={}", psi.collapse(0));

    tebd::runBrickwall(psi, bonds, chiMax, eps, nSteps, fUpdate);
    // std::println("E = {}", model.energy(psi).real());
    // std::println("<N> = {}", psi.getSiteExpectationValueSum(particleNumberOperator).real() / model.L);
    // std::println("m = {}", psi.getSiteExpectationValues(sigma_z));
    // std::println("chiS: {:b}", psi.getChi());
    // psi.print();
    // std::println("<psi|psi>={}", psi.collapse(0));
}


void testTFI_TEBD() {
    unsigned L = 40;
    std::println("Make model");
    auto model = TransverseFieldIsingModel1D<dtype>(L, dtype(1.), dtype(1.5));
    // auto& mpo = *(model.getH_MPO().front());
    // e::Tensor<std::complex<float>, 4> mpo(3, 3, 2, 2);
    // mpo.setValues({
    //     {
    //         {{0,1},{2,3}},
    //         {{4,5},{6,7}},
    //         {{8,9},{10,11}}
    //     }, {
    //         {{12,13},{14,15}},
    //         {{16,17},{18,19}},
    //         {{20,21},{22,23}},
    //     }, {
    //         {{24,25},{26,27}},
    //         {{28,29},{30,31}},
    //         {{32,33},{34,35}}

    //     }});
    // std::cout << "first MPO: \n" << mpo  << std::endl;
    // std::cout << "mapped : \n" << e::Map<e::Matrix<dtype, 6, 6, e::RowMajor>>(mpo.data()) << std::endl;
    // for (unsigned i = 0; i < mpo.size(); i++) {
    //     std::cout << mpo.data()[i] << " ";
    //     if ((i+1) % 2 == 0) { std::cout << "\n"; }
    //     if ((i+1) % 4 == 0) { std::cout << "\n"; }
    // }
    e::Tensor<dtype, 2> sigma_z(2, 2);
    sigma_z.setValues({{1, 0}, {0, -1}});

    dtype eps = 1e-20;
    unsigned chiMax = 30;
    dtype dt = 0.05;
    std::println("dt={}", dt);
    std::cout << "leftT:\n" << *(model.getH_Bonds().front()) << std::endl;
    std::cout << "leftT2:\n" << bondMatrix2Tensor(*(model.getH_BondsMatrices().front()), model.localDim) << std::endl;
    std::cout << "middleT:\n" << *(model.getH_Bonds().at(1)) << std::endl;
    auto bonds = tebd::getTimeEvolutionMatrices(model.getH_BondsMatrices(), dt, model.localDim);
    std::cout << "U_Bond:\n" << *(bonds.front()) << std::endl;
    // std::cout << "H_Bond:\n" << *(bonds.front()) << std::endl;

    // MPS<dtype> psi = initMPS_spinUp<dtype>(L);
    MPS<dtype> psi = initMPS_spinRight<dtype>(L);

    unsigned nSteps = 200;
    postUpdateFunction_t<dtype> fUpdate = postUpdateNoop<dtype>;
    // auto fUpdate = postUpdateFunction_t<dtype>([](const MPS<dtype>& psi, unsigned i){
    //     std::println("i={:03}: <psi|psi>={}", i, psi.collapse(0));
    // });
    std::println("E = {}", model.energy(psi).real());
    std::println("m = {}", psi.getSiteExpectationValues(sigma_z));
    std::println("chiS: {}", psi.getChi());
    psi.print();
    std::println("<psi|psi>={}", psi.collapse(0));

    tebd::runBrickwall(psi, bonds, chiMax, eps, nSteps, fUpdate);
    // updateBond(psi, 0, *(bonds.at(0)), chiMax, eps);
    // updateBond(psi, 2, *(bonds.at(2)), chiMax, eps);
    // updateBond(psi, 4, *(bonds.at(4)), chiMax, eps);
    // updateBond(psi, 6, *(bonds.at(6)), chiMax, eps);
    // updateBond(psi, 1, *(bonds.at(1)), chiMax, eps);
    // updateBond(psi, 3, *(bonds.at(3)), chiMax, eps);
    // updateBond(psi, 5, *(bonds.at(5)), chiMax, eps);
    // updateBond(psi, 7, *(bonds.at(7)), chiMax, eps);


    std::println("E = {}", model.energy(psi).real());
    std::println("m = {}", psi.getSiteExpectationValues(sigma_z));
    std::println("chiS: {}", psi.getChi());
    // psi.print();
    std::println("<psi|psi>={}", psi.collapse(0));

    // for (unsigned i = 0; i < L; i++) {
    //     std::println("i={:02}: <psi|psi>={}",i, psi.collapse(i));
    // }
}


void testTFI_DMRG() {
    unsigned L = 8;
    std::println("Make model");
    auto model = TransverseFieldIsingModel1D<dtype>(L, dtype(1.), dtype(1));
    e::Tensor<dtype, 2> sigma_z(2, 2);
    sigma_z.setValues({{1, 0}, {0, -1}});

    std::cout << *(model.getH_BondsMatrices().at(0)) << std::endl;
    printTensor(*model.getH_Bonds().at(0), "H0");
    std::cout << *(model.getH_BondsMatrices().at(1)) << std::endl;
    printTensor(*model.getH_Bonds().at(1), "H1");
    std::cout << *(model.getH_BondsMatrices().at(6)) << std::endl;
    printTensor(*model.getH_Bonds().at(6), "H6");


    dtype eps = 1e-20;
    unsigned chiMax = 10;
    MPS<dtype> psi = initMPS_spinUp<dtype>(L);
    for (unsigned i = 0; i < 7; i++) {
        std::println("{}", psi.getBondExpectationValue(*(model.getH_Bonds()[i]), i));
    }
    // MPS<dtype> psi = initMPS_spinRight<dtype>(L);
    // MPS<dtype> psi = initMPS_spinDown<dtype>(L);
    dtype E = model.energy(psi);
    std::println("E{} = {}", -1, E.real());
    std::println("M{} = {}", -1, psi.getSiteExpectationValueSum(sigma_z));

    printTensor(*(model.getH_MPO().at(0)), "H_mpo[0]");
    printTensor(*(model.getH_MPO().at(1)), "H_mpo[1]");

    auto engine = dmrg::DMRGEngine(psi, model.getH_MPO(), chiMax, eps);
    std::println("Initialized");
    E = model.energy(psi);
    std::println("E{} = {}", 0, E);
    // engine.update_bond(0);
    for (unsigned i = 0; i < 4; i++) {
        engine.sweep(psi);
        std::println("\nAfter sweep #{}", i+1);
        E = model.energy(psi);
        std::println("E = {:.2f}", E);
        std::println("M = {:.2f}", psi.getSiteExpectationValueSum(sigma_z));
    }

    // engine.psi.print();

}

int main() {
    testTFI_DMRG();
    // testSVD(); return 0;
    // testBoseHubbardTEBD();
    // testTFI_TEBD();
    // testWasmCode();
    // return runBenchmark_oneAtATime();
    // return runBenchmark_all();

    // std::println("Start calculations");
    // e::Tensor<dtype, 3> t1(10, 2, 10);
    // t1.setZero();
    // std::cout << t1 << std::endl;
    // t1(0,1,0) = .69;
    // std::cout << t1 << std::endl;

    // e::Tensor<dtype, 3> t2(10, 2, 10);
    // t2.setRandom();
    // std::cout << t2 << std::endl;

    // e::array<iP, 3> productDims = { iP(0, 0), iP(1, 1), iP(2, 2) };
    // auto t3 = t1.contract(t2, productDims);
    // std::cout << t3 << std::endl;

    // double eps = 1e-8;
    // size_t chi_max = 5;   

    // verified energy function works as expected
    // unsigned L = 14;
    // for (double g : {0.5, 1., 1.5}) {
    //     auto model2 = TFIModel<dtype>(L, dtype(1.), dtype(g));
    //     auto mps = initMPSspinUp<dtype>(L);
    //     std::println("g={} => E = {}", g, model2.energy(mps));
    //
    return 0;
}
