#include "dmrg.hpp"

#include "algorithm.hpp"
#include "type.hpp"
#include "util.hpp"


#include <lambda_lanczos/lambda_lanczos.hpp>
using lambda_lanczos::LambdaLanczos;

#include <print>
#include <vector>
#include <memory>
#include <cmath>

namespace dmrg {
	void HeffTheta(
            const e::Tensor<dtype, 3>& LP,  // vL wL* vL*
            const e::Tensor<dtype, 3>& RP,  // vR* wR* vR
            const e::Tensor<dtype, 4>& W1,  // wL wC i i*
            const e::Tensor<dtype, 4>& W2,  // wC wR j j*
            const std::vector<dtype>& thetaIn,
            std::vector<dtype>& thetaOut
    ) {
        assert(thetaIn.size() == LP.dimension(0) * W1.dimension(2) * W2.dimension(2) * RP.dimension(2));
        assert(thetaOut.size() == LP.dimension(0) * W1.dimension(2) * W2.dimension(2) * RP.dimension(2));
        e::array<e::Index, 4> thetaShape = {LP.dimension(0), W1.dimension(2), W2.dimension(2), RP.dimension(2)};
        // Reshape theta into a 4D tensor
        e::TensorMap<const e::Tensor<dtype, 4>> thetaTensor(thetaIn.data(), thetaShape);    // vL i j vR
        // e::Tensor<dtype, 4> x = theta_tensor;
        auto LPT = LP.contract(thetaTensor, e::array<iP, 1>({iP{2, 0}}));                   // vL wL* [vL*], [vL] i j vR
        auto LPTW1 = LPT.contract(W1, e::array<iP, 2>({iP{1, 0}, iP{2, 3}}));               // vL [wL*] [i] j vR, [wL] wC i [i*]
        auto LPTW1W2 = LPTW1.contract(W2, e::array<iP, 2>({iP{3, 0}, iP{1, 3}}));           // vL [j] vR [wC] i, [wC] wR j [j*]
        auto LPTW1W2RP = LPTW1W2.contract(RP, e::array<iP, 2>({iP{1, 0}, iP{3, 1}}));       // vL [vR] i [wR] j, [vR*] [wR*] vR

        // assigning to map will assign to underlying std::vector
        e::TensorMap<e::Tensor<dtype, 4>> thetaOutMap(thetaOut.data(), thetaShape);         // vL i j vR
        thetaOutMap = LPTW1W2RP;
        // std::println("thetaO: {}", thetaOut);
    }


    template<typename T>
    DMRGEngine<T>::DMRGEngine(MPS<T>& psi, const std::vector<std::shared_ptr<e::Tensor<T, 4>>>& H_mpo,
               unsigned chiMax, T eps)
        : H_mpo(H_mpo), chiMax(chiMax), eps(eps) {
        int L = psi.getL();
        LPs.resize(L);
        RPs.resize(L);
        initialize(psi);
    }

	template<typename T>
    void DMRGEngine<T>::initialize(MPS<T>& psi) {
        // Initialize LPs and RPs
        int D = H_mpo[0]->dimension(0);
        int chi = psi.B(0).dimension(0);
        e::Tensor<T, 3> LP(chi, D, chi), RP(chi, D, chi);
        LP.setZero();
        RP.setZero();
        LP.chip(0, 1) = identityTensor<T>(chi, chi);
        RP.chip(D - 1, 1) = identityTensor<T>(chi, chi);

        LPs[0] = LP;
        RPs.back() = RP;

        for (int i = psi.getL() - 1; i > 1; --i) {
            update_RP(i, psi);
        }
    }

	template<typename T>
    void DMRGEngine<T>::sweep(MPS<T>& psi, postUpdateFunction_t<T> postUpdateFunc) {
        // Sweep left to right
        for (unsigned i = 0; i < psi.getL() - 2; ++i) {
            update_bond(i, psi);
            postUpdateFunc(psi, i);
        }
        // Sweep right to left
        for (unsigned i = psi.getL() - 2; i > 0; --i) {
            update_bond(i, psi);
            postUpdateFunc(psi, i);
        }
    }

	template<typename T>
    void DMRGEngine<T>::update_bond(int i, MPS<T>& psi) {
        int j = i + 1;
        // std::println("update({})", i);

        e::Index d1   = H_mpo[i]->dimension(2);
        e::Index d2   = H_mpo[j]->dimension(2);
        e::Index chi1 = LPs[i].dimension(0);
        e::Index chi2 = RPs[j].dimension(2);
        e::Index dim  = chi1*d1*d2*chi2;

        // 1) Find GS (smallest eigenvalue / vector)
        // Approaches 1 and 2 are wrong, because we need a more complicated contraction series
        // instead of a reshape + simple matrix product for Heff @ theta

        // For i) and ii): create effective hamiltonian
        // auto LW = LPs[i].contract(*H_mpo[i], e::array<iP, 1>{iP{1, 0}}); // vL [wL] vL*,  [wL] wR i i* 
        // auto LWW = LW.contract(*H_mpo[j], e::array<iP, 1>{iP{2, 1}});    // vL vL* [wR] i i*, [wL] wR j j*
        // auto LWWR = LWW.contract(RPs[j], e::array<iP, 1>{iP{4, 1}});     // vL vL* i i* [wR] j j*, vR* [wR*] vR
        // e::Tensor<T, 8> Heff = LWWR;// vL vL* i i* j j*, vR* vR
        // printTensorDims(Heff, "Heff");

        // printTensor(LPs[i], "LP[i]");
        // printTensor(RPs[j], "RP[j]");
        // printTensor(*H_mpo[i], "H_mpo[i]");
        // printTensor(*H_mpo[j], "H_mpo[j]");
        // e::Map<e::MatrixX<T>> HeffMat(Heff.data(), dim, dim);  // vL vL* i i*, j j* vR* vR
        // std::cout << "HeffMat" << std::endl;
        // std::cout << HeffMat << std::endl;

        // i) Eigen
        // Realllyyyy inefficient, it calculates all eigenvalues and we just need the eigenvector to the
        // smallest absolute eigenvalue
        // e::SelfAdjointEigenSolver<e::MatrixX<T>> es(HeffMat);
        // const e::VectorX<T>& S = es.eigenvalues();
        // const e::MatrixX<T>& U = es.eigenvectors();
        // std::println("Eigenvalues with eigen:");
        // std::cout << S << std::endl;
        // std::println("Eigenvector with eigen:");
        // std::cout << U.rows() << "x" << U.cols() << std::endl;
        // e::Map<const e::MatrixX<T>> U1(U.col(0).data(), chi1*d1, d2*chi2);
        // std::cout << "theta2 (eigen)" << std::endl;
        // std::cout << U1 << std::endl;
        // find the smallest eigenvalue
        // e::TensorMap<const e::Tensor<T, 4>> theta(U.data(), chi1, d1, d2, chi2);
        // auto [Ai, Sj, Bj] = split_truncate_theta<T>(theta, chiMax, eps);
        // auto [Ai, Sj, Bj] = split_truncate_theta<T>(U1, chi1, d1, d2, chi2, chiMax, eps);

        // ii) Lanczos
        // std::println("Eigenvector with Lanczos:");
        // auto [eval, theta2Vec] = lanczos(HeffMat, psi.getTheta2(i));
        // e::Map<e::MatrixX<T>> theta2(theta2Vec.data(), chi1*d1, d2*chi2);
        // std::cout << theta2.rows() << "x" << theta2.cols() << std::endl;
        // std::cout << eval << std::endl;
        // std::cout << theta2 << std::endl;
        // auto [Ai, Sj, Bj] = split_truncate_theta<T>(theta2, chi1, d1, d2, chi2, chiMax, eps);
        // auto theta0 = psi.getTheta2(i);
        // printTensor(theta0, "theta0");
        // e::Map<e::VectorX<T>> theta0v(theta0.data(), theta0.size());
        // std::println("Heff @ theta0");
        // e::VectorX<T> v = HeffMat * theta0v;
        // std::cout << v << std::endl;

        // iii) Lanczos with custom contraction function
        e::Tensor<T, 4> theta0 = psi.getTheta2(i);
        // Use Theta2 as the initial guess for the GS
        auto initF = [&theta0](std::vector<T>& v) {
            v.resize(theta0.size());
            for (unsigned i = 0; i < theta0.size(); i++) {
                v.at(i) = theta0.data()[i];  // TODO find nicer solution
            }
            // v = std::vector<T>(theta0.data(), theta0.size());
        };
        auto matmul = [this, i, j](const std::vector<T>& thetaIn, std::vector<T>& thetaOut) {
            HeffTheta(this->LPs.at(i), this->RPs.at(j), *(this->H_mpo.at(i)), *(this->H_mpo.at(j)), thetaIn, thetaOut);
        };

        // TODO set eps
        LambdaLanczos<T> engine(matmul, dim, false, 1); // Find 1 maximum eigenvalue
        engine.init_vector = initF;
        std::vector<double> eigenvalues;
        std::vector<std::vector<T>> eigenvectors;
        engine.run(eigenvalues, eigenvectors);
        e::Map<e::MatrixX<T>> theta2(eigenvectors.front().data(), chi1*d1, d2*chi2);
        // std::println("Lanczos ev: {}", eigenvalues);
        // e::TensorMap<e::Tensor<T, 4>> theta2T(eigenvectors.front().data(), chi1, d1, d2, chi2);
        // printTensor(theta2T, "theta2");

        auto [Ai, Sj, Bj] = split_truncate_theta<T>(theta2, chi1, d1, d2, chi2, chiMax, eps);



        // 2) Truncate
        // printTensor(Ai, "Ai");
        // printTensor(Sj, "Sj");
        // printTensor(Bj, "Bj");
        // std::cout << "a: " << a.dimensions() << "\nA: " << Ai.dimensions() << "\nS: " << Sj.dimensions() << "\nB: " << Bj.dimensions() << std::endl;
        // renormalize to Bi
        e::Tensor<T, 1> SiInv = 1. / psi.S(i);
        e::Tensor<T, 2> SiDiv = diagonal(SiInv);
        e::Tensor<T, 3> Gi = SiDiv.contract(Ai, e::array<iP, 1>{iP{1, 0}}); // vL [vL], [vL] i vC
        // printTensor(SiDiv, "SiDiv");
        // printTensor(Gi, "Gi");

        // e::Tensor<T, 0> SS = Sj.contract(Sj.conjugate(), e::array<iP, 1>{iP{0,0}});
        // std::println("Bond {:03}, Sj contracted with its conjugate = {}", i, SS(0));

        // psi.B(i) = Gi.contract(diagonal(Sj), e::array<iP, 1>{iP{2, 0}}); // vL i [vC], [vC] vC
        // psi.B(i+1) = std::move(Bj);
        // psi.S(i+1) = std::move(Sj);
        typename MPS<T>::BTensor_t Bi = Gi.contract(diagonal(Sj), e::array<iP, 1>{iP{2, 0}}); // vL i [vC], [vC] vC
        // e::Tensor<bool, 0> equal = (psi.B(i) == Bi).all();
        // printTensor(psi.B(i), "old Bi");
        // printTensor(Bi, "new Bi");
        psi.B(i) = std::move(Bi);
        psi.B(j) = std::move(Bj);
        psi.S(j) = std::move(Sj);

        // 3) Update
        update_LP(i, psi);
        update_RP(j, psi);
    }

	template<typename T>
    void DMRGEngine<T>::update_RP(int i, MPS<T>& psi) {
        // std::println("update_RP({})", i);
        int j = i - 1;
        const e::Tensor<T, 3>& RP = RPs[i];            // vR* wR* vR
        const e::Tensor<T, 3>& B = psi.B(i);           // vL i vR
        e::Tensor<T, 3> Bc = B.conjugate();     // vL* i* vR*
        const e::Tensor<T, 4>& W = *H_mpo[i];          // wL wR i i*
#ifdef DMRGDEBUG
        printTensorDims(RP, "RP");
        printTensorDims(B, "B");
        printTensorDims(Bc, "Bc");
        printTensorDims(W, "W");

        auto RP2  = B.contract(RP, e::array<iP, 1>{iP{2, 0}});            // vL i [vR], [vR*] wR* vR
        e::Tensor<T, 4> RP2T = RP2;
        printTensorDims(RP2T, "RP2T");
        auto RP3 = RP2T.contract(W, e::array<iP, 2>{iP{1, 3}, iP{2, 1}}); // vL [i] [wR*] vR, wL [wR] i [i*]
        e::Tensor<T, 4> RP3T = RP3;
        printTensorDims(RP3T, "RP3T");
        auto RP4 = RP3.contract(Bc, e::array<iP, 2>{iP{1, 2}, iP{3, 1}}); // vL [vR] wL [i], vL* [i*] [vR*]
        e::Tensor<T, 3> RP4T = RP4;
        printTensor(RP4T, "RP4T");
#else
        auto RP2 = B.contract(RP, e::array<iP, 1>{iP{2, 0}});               // vL i [vR], [vR*] wR* vR
        auto RP3 = RP2.contract(W, e::array<iP, 2>{iP{1, 3}, iP{2, 1}});    // vL [i] [wR*] vR, wL [wR] i [i*]
        auto RP4 = RP3.contract(Bc, e::array<iP, 2>{iP{1, 2}, iP{3, 1}});   // vL [vR] wL [i], vL* [i*] [vR*]
#endif
        RPs[j] = RP4;
    }

	template<typename T>
    void DMRGEngine<T>::update_LP(int i, MPS<T>& psi) {
        // std::println("update_LP({})", i);
        int j = i + 1;
        const e::Tensor<T, 3>& LP = LPs[i];    // vL wL vL*
        const e::Tensor<T, 3>& B = psi.B(i);   // vL i vR
        const e::Tensor<T, 4>& W = *H_mpo[i];  // wL wR i i*

        e::Tensor<T, 1> SjInvV = 1. / psi.S(j);
        e::Tensor<T, 2> SjInv = diagonal(SjInvV);
        auto G = B.contract(SjInv, e::array<iP, 1>{iP{2, 0}});    // vL i [vR], [vR*] vR
        e::Tensor<T, 2> Si = diagonal(psi.S(i));
        auto A = Si.contract(G, e::array<iP, 1>{iP{1, 0}});       // vL [vL*], [vL] i vR
        auto Ac = A.conjugate();            // vL* i* vR*

#ifdef DMRGDEBUG
        auto LP2 = LP.contract(A, e::array<iP, 1>{iP{2, 0}});     // vL wL* [vL*], [vL] i vR
        e::Tensor<T, 4> LP2T = LP2;
        printTensorDims(LP2T, "LP2T");
        printTensorDims(W, "W");
        auto LP3 = W.contract(LP2T, e::array<iP, 2>{iP{0, 1}, iP{3, 2}});  // [wL] wR i [i*], vL [wL*] [i] vR
        e::Tensor<T, 4> LP3T = LP3;
        printTensorDims(LP3T, "LP3T");
        e::Tensor<T, 3> AcT = Ac;
        printTensorDims(AcT, "AcT");
        auto LP4 = AcT.contract(LP3, e::array<iP, 2>{iP{0, 2}, iP{1, 1}}); // [vL*] [i*] vR*, wR [i] [vL] vR
        e::Tensor<T, 3> LP4T = LP4;
        printTensor(LP4T, "LP4T");
        LPs[j] = LP4T;
#else
        auto LP2 = LP.contract(A, e::array<iP, 1>{iP{2, 0}});               // vL wL* [vL*], [vL] i vR
        auto LP3 = W.contract(LP2, e::array<iP, 2>{iP{0, 1}, iP{3, 2}});    // [wL] wR i [i*], vL [wL*] [i] vR
        auto LP4 = Ac.contract(LP3, e::array<iP, 2>{iP{0, 2}, iP{1, 1}});   // [vL*] [i*] vR*, wR [i] [vL] vR
        LPs[j] = LP4;
#endif
    }

    template<typename T>
    void run(MPS<T>& psi, const std::vector<std::shared_ptr<e::Tensor<T, 4>>>& H_mpos, unsigned chiMax, T eps, unsigned nSweeps, postUpdateFunction_t<T> postSweepFunc, postUpdateFunction_t<T> postUpdateFunc) {
        auto engine = DMRGEngine(psi, H_mpos, chiMax, eps);
        for (unsigned i = 0; i < nSweeps; i++) {
            engine.sweep(psi, postUpdateFunc);
            if (postSweepFunc(psi, i)) break;
        }
    }

    template class dmrg::DMRGEngine<dtype>;
	template void run<dtype>(MPS<dtype>& psi, const std::vector<std::shared_ptr<e::Tensor<dtype, 4>>>& H_mpos, unsigned chiMax, dtype eps, unsigned nSweeps, postUpdateFunction_t<dtype>, postUpdateFunction_t<dtype>);

}  // namespace dmrg

