#include "tebd.hpp"
#include "type.hpp"
#include "util.hpp"
#include "model.hpp"
#include <cstdio>
#include <unordered_map>
#include <print>
#include <eigen3/unsupported/Eigen/MatrixFunctions>

template<typename InType, typename OutType>
std::vector<std::shared_ptr<OutType>> forEachUnique(std::function<OutType(const InType&)> f, const std::vector<std::shared_ptr<InType>>& inRange) {
    std::unordered_map<std::shared_ptr<InType>, std::shared_ptr<OutType>> cache;
    std::vector<std::shared_ptr<OutType>> outRange;
    outRange.reserve(inRange.size());
    for (const auto& inPtr : inRange) {
        if (cache.contains(inPtr)) {
            outRange.push_back(cache.at(inPtr));
        } else {
            std::shared_ptr<OutType> outPtr = std::make_shared<OutType>(f(*inPtr));
            // cache.emplace(std::make_pair(inPtr, outPtr));
            cache.emplace(inPtr, outPtr);
            outRange.push_back(outPtr);
        }
    }
    return outRange;
}


namespace tebd {
    template<typename T>
    e::Tensor<T, 4> getTimeEvolutionMatrix(const e::MatrixX<T>& bond, T dt, unsigned localDim) {
        e::MatrixX<T> U = (bond * (Imag*dt)).exp();
        // std::cout << "U = \n" << U << "\n\n";
        return bondMatrix2Tensor<T>(U, localDim);
    }

    template<typename T>
    Bonds<T> getTimeEvolutionMatrices(const std::vector<std::shared_ptr<e::MatrixX<T>>>& bonds, T dt, unsigned int localDim) {
        std::function<e::Tensor<T, 4>(const e::MatrixX<T>&)> f = std::bind(getTimeEvolutionMatrix<T>, std::placeholders::_1, dt, localDim);
        std::vector<std::shared_ptr<e::Tensor<T, 4>>> U_Bonds = forEachUnique(f, bonds);
        return U_Bonds;
    }


    template<typename T>
    void updateBond(MPS<T>& psi, size_t i, const e::Tensor<T, 4>& U, unsigned chiMax, T eps) {
        // std::println("BOND {}", i);
        auto theta = psi.getTheta2(i);  // vL i j vR
        // apply U: i j i* j*
        // contraction gives vL vR i j -> shuffle to vL i j vR
        // std::cout << theta.dimension(2) << " - " << U.dimension(2) << " / "  << theta.dimensions() << std::endl;
        e::Tensor<T, 4> a = theta.contract(U, e::array<iP, 2>({iP{1, 2}, iP{2, 3}}));
        e::Tensor<T, 4> b = a.shuffle(e::array<int, 4>{0, 2, 3, 1});
        // e::Tensor<T, 4> b = a.shuffle(e::array<int, 4>{2, 0, 1, 3});
        // printTensor(b, "Utheta");
        // SVD and truncation
        auto [Ai, Sj, Bj] = split_truncate_theta<T>(b, chiMax, eps);
        // 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);
        MPS<dtype>::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) = Bi;
        psi.B(i+1) = Bj;
        psi.S(i+1) = Sj;
    }


    template<typename T>
    void runBrickwall(MPS<T>& psi, const Bonds<T>& bonds, unsigned chiMax, T eps, unsigned nSteps, postUpdateFunction_t<T> postUpdateFunc) {
        for (unsigned n = 0; n < nSteps; n++) {
            // std::print("i = {:03}\r", n);
            // std::flush(std::cout);
            for (unsigned row = 0; row < 2; row++) {
                for (unsigned i = row; i < bonds.size(); i+=2) {
                    // std::cout << "i = " << i << std::endl;
                    // std::println("Updating bond {:02}", i);
                    updateBond(psi, i, *(bonds.at(i)), chiMax, eps);
                }
            }
            if (postUpdateFunc(psi, n)) break;
        }
    }

    // TODO
    /**
     * @brief Run TEBD using a second-order decomposition
     * @todo I think this is currently only correct for even system sizes => check decomposition
     */
    template<typename T>
    void runSecondOrder(MPS<T>& psi, const Bonds<T>& bonds, const Bonds<T>& bondsHalfDT, unsigned chiMax, T eps, unsigned nSteps, postUpdateFunction_t<T> postUpdateFunc) {
        for (unsigned n = 0; n < nSteps; n++) {
            // std::print("i = {:03}\r", n);
            // std::flush(std::cout);
            for (unsigned row = 0; row < 2; row++) {
                const Bonds<T>* useBonds = &bonds;
                if (row == 0 && (n != 0 || n != nSteps-1)) {
                    useBonds = &bondsHalfDT;
                }
                for (unsigned i = row; i < bonds.size(); i+=2) {
                    // std::cout << "i = " << i << std::endl;
                    updateBond(psi, i, *(useBonds->at(i)), chiMax, eps);
                }
            }
            if (postUpdateFunc(psi, n)) break;
        }
    }


    // Instantiations
    template e::Tensor<dtype, 4> getTimeEvolutionMatrix<dtype>(const e::MatrixX<dtype>& bond, dtype dt, unsigned localDim);
    template Bonds<dtype> getTimeEvolutionMatrices<dtype>(const std::vector<std::shared_ptr<e::MatrixX<dtype>>>& bonds, dtype dt, unsigned int localDim);
    template void updateBond<dtype>(MPS<dtype>& psi, size_t i, const e::Tensor<dtype, 4>& U, unsigned chiMax, dtype eps);
    template void runBrickwall<dtype>(MPS<dtype>& psi, const Bonds<dtype>& bonds, unsigned chiMax, dtype eps, unsigned nSteps, postUpdateFunction_t<dtype> postUpdateFunc);
    template void runSecondOrder<dtype>(MPS<dtype>& psi, const Bonds<dtype>& bonds, const Bonds<dtype>&, unsigned chiMax, dtype eps, unsigned nSteps, postUpdateFunction_t<dtype> postUpdateFunc);

};  // namespace tebd
template std::vector<std::shared_ptr<e::Tensor<dtype, 4>>> forEachUnique(std::function<e::Tensor<dtype, 4> (const e::MatrixX<dtype> &)> f, const std::vector<std::shared_ptr<e::MatrixX<dtype>>>&);
