#pragma once
#include "algorithm.hpp"
#include "mps.hpp"
#include <eigen3/Eigen/Eigen>


/**
 * @brief Returns a \f$L\times L \f$ identity matrix in tensor form, where \f$L = \text{min}(N, M)\f$
 */
template <typename T>
e::Tensor<T, 2> identityTensor(e::Index N, e::Index M) {
    e::Tensor<dtype, 2> t(N, M);
    t.setZero();
    e::Index min = std::min(N, M);
    for (e::Index i = 0; i < min; i++) {
        t(i, i) = 1;
    }
    return t;
}


/**
 * @brief Implementation of the Density Matrix Renormalization Group (DMRG) algorithm
 */
namespace dmrg {
    /**
     * @brief Helper class that keeps the left and right blocks of the MPS
     * @note Always call with the same mps (psi) that it was initialized with!
     *  It would have been better if the engine would take ownership of the MPS during sweeps,
     *  but this would make calling an update function in between more difficult.
     */
    template<typename T>
    class DMRGEngine {
    public:
        std::vector<std::shared_ptr<e::Tensor<T, 4>>> H_mpo;
        std::vector<e::Tensor<T, 3>> LPs, RPs;
        int chiMax;
        T eps;

        DMRGEngine(MPS<T>& psi, const std::vector<std::shared_ptr<e::Tensor<T, 4>>>& H_mpo,
                   unsigned chiMax = 100, T eps = 1e-12);

        void initialize(MPS<T>& psi); 
        /// Post update function is called after each bond update, its only meant for emscripten_sleep!
        void sweep(MPS<T>& psi, postUpdateFunction_t<T> postUpdateFunc=postUpdateNoop<T>);
        void update_bond(int i, MPS<T>& psi);
        void update_RP(int i, MPS<T>& psi); 
        void update_LP(int i, MPS<T>& psi);
    };

    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=postUpdateNoop<T>);
}
