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

/**
 * @brief Apply the function f on every element if inRange, but reuse a cached result when two pointers point to the same object
 */
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);

template<typename T>
using Bonds = std::vector<std::shared_ptr<e::Tensor<T, 4>>>;

/**
 * @brief Implementation of the Time Evolved block Decimation (TEBD) algorithm
 */
namespace tebd {
    /**
     * @brief Apply time evolution on a bond
     * @details
     *  \f[U = \exp{ B \, i \partial t} \f]
     */
    template<typename T>
    e::Tensor<T, 4> getTimeEvolutionMatrix(const e::MatrixX<T>& bond, T dt, unsigned localDim);

    template<typename T>
    Bonds<T> getTimeEvolutionMatrices(const std::vector<std::shared_ptr<e::MatrixX<T>>>& bonds, T dt, unsigned int localDim);

    /**
     * @brief Apply U on bond i, perform SVD and truncate, update MPS
     */
    template<typename T>
    void updateBond(MPS<T>& psi, size_t i, const e::Tensor<T, 4>& U, unsigned chiMax, T eps);


    /**
     * @brief Brickwall TEBS scheme
     * @details
     *  Apply U on every second bond starting from the first,
     *  then apply U on every second bond starting from the second
     *
     *  First order scheme
     *
     * @param postUpdateFunc function that can be used to track expectation values after each update
     */
    template<typename T>
    void runBrickwall(MPS<T>& psi, const Bonds<T>& bonds, unsigned chiMax, T eps, unsigned nSteps, postUpdateFunction_t<T> postUpdateFunc);

    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);
};  // namespace tebd
