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

/**
 * @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 Apply time evolution on a bond
 * @details
 *  \f[U = \exp{ B \, \del 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);

template<typename T>
using postUpdateFunction_t = std::function<void (const MPS<T>&, unsigned)>;

template<typename T>
void postUpdateNoop(const MPS<T>&, unsigned) {};

/**
 * @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 runTEBD_brickwall(MPS<T>& psi, const Bonds<T>& bonds, unsigned chiMax, T eps, unsigned nSteps, postUpdateFunction_t<T> postUpdateFunc);

template<typename T>
void runTEBD_secondOrder(MPS<T>& psi, const Bonds<T>& bonds, const Bonds<T>& bondsHalfDT, unsigned chiMax, T eps, unsigned nSteps, postUpdateFunction_t<T> postUpdateFunc);