#pragma once
#include "util.hpp"
#include <cmath>
#include <eigen3/Eigen/Eigen>
#include <eigen3/unsupported/Eigen/CXX11/Tensor>
#include <vector>
#include <cassert>
#include <random>

#include <iostream>

namespace e = Eigen;
using iP = e::IndexPair<int>;

template<typename T>
e::Tensor<T, 2> diagonal(const e::Tensor<T, 1>& t) {
    e::Tensor<T, 2> out(t.dimension(0), t.dimension(0));
    out.setZero();
    for (e::Index i = 0; i < t.dimension(0); i++) {
        out(i, i) = t(i);
    }
    return out;
}

using namespace std::literals;

/**
 * @brief Matrix Product State class
 * @todo Possible Optimization: Store S as full tensor and not just the diagonal values: we always need to construct a diagonal tensor with `diagonal` function which leads to memory allocations. Shouldnt be too bad from storage perspective
 */
template<typename T>
class MPS {
    public:
        using BTensor_t = e::Tensor<T, 3>;
        using STensor_t = e::Tensor<T, 1>;
        MPS() {};
        MPS(std::vector<BTensor_t>&& Bs, std::vector<STensor_t>&& Ss);
        /**
         * @brief Return single site contraction of the $S_i * B_i$ tensors
         * @returns Contraction with legs `vL i vR`
         */
        BTensor_t getTheta1(size_t siteIdx) const;

        /**
         * @brief Return two-site contraction of the $S_i * S_B * B_{i+1}$ tensors
         * @returns Contraction with legs `vL i i+1 vR`
         */
        e::Tensor<T, 4> getTheta2(size_t bondIdx) const;

        /**
         * @brief Return the dimensions of all bonds
         */
        std::vector<unsigned> getChi() const;
        /// @brief Return the maximum bond dimension
        unsigned getChiMax() const;
        /// @brief Return the average bond dimension
        double getChiAverage() const;
        
        /**
         * @brief Compute the expectation value of a single-site operator on a single site
         * @param op Single-site operator with legs i i*
         */
        T getSiteExpectationValue(const e::Tensor<T, 2>& op, size_t siteIdx) const;
        std::vector<T> getSiteExpectationValues(const e::Tensor<T, 2>& op) const;
        T getSiteExpectationValueSum(const e::Tensor<T, 2>& op) const;

        /**
         * @brief Compute the expectation value of a two-site operator on a single bond
         * @param op Two-site operator with legs i j i* j*
         */
        T getBondExpectationValue(const e::Tensor<T, 4>& op, size_t bondIdx) const;

        T getEntanglementEntropy(size_t bondIndex) const;

        /**
         * @brief Collapse the MPS, calculating \f[<\psi|\psi>\f]
         */
        T collapse(size_t sizeIndex) const;

        BTensor_t& B(size_t i) { return this->Bs.at(i); }
        STensor_t& S(size_t i) { return this->Ss.at(i); }
        const BTensor_t& B(size_t i) const { return this->Bs.at(i); }
        const STensor_t& S(size_t i) const { return this->Ss.at(i); }
        // size_t getPhysicalDimension() const { return this->Bs.at(0).dimension(1); }
        size_t getL() const { return this->L; }
        void print() const {
            for (size_t i = 0; i < this->L; i++) {
                printTensor(Ss.at(i), "S["s + std::to_string(i) + "]");
                printTensor(Bs.at(i), "B["s + std::to_string(i) + "]");
            }
        }
    private:
        std::vector<BTensor_t> Bs;
        std::vector<STensor_t> Ss;
        size_t L;
};


/**
 * @brief Function to perform split and truncate on a two-site wave function
 * @param theta Two-site wave function in mixed canonical form, with legs `vL, i, j, vR`
 * @param chi_max Maximum number of singular values to keep
 * @param eps Discard any singular values smaller than that.
 *
 * @returns (A, S, B) where A is the left-canonical matrix on site i (legs `vL, i, vC`),
 *     S are the singular values and
 *     B is the right-canonical matrix on site i+1, with legs `vC, j, vR`
 * @note This functions maps the tensor and calls the Matrux overload of split_truncate_theta internally.
 */
template<typename dtype>
std::tuple<e::Tensor<dtype, 3>, e::Tensor<dtype, 1>, e::Tensor<dtype, 3>> 
split_truncate_theta(const e::Tensor<dtype, 4>& theta, size_t chiMax, dtype eps);

template<typename dtype>
std::tuple<e::Tensor<dtype, 3>, e::Tensor<dtype, 1>, e::Tensor<dtype, 3>> 
split_truncate_theta(const e::MatrixX<dtype>& theta, e::Index chivL, e::Index dL, e::Index dR, e::Index chivR, size_t chiMax, dtype eps);


/**
 * @brief Apply a pertubation to the entries in the physical dimension of the B tensors
 * @details
 *  For each value in the physical dimension, adds a random value between 0 and `pertubationStrength`.
 *  Then, every tensor is normalized again.
 */
template<typename T>
void applyPertubationOnB(std::vector<e::Tensor<T, 3>>& Bs, double pertubationStrength) {
    std::mt19937 generator(493852934);  // Mersenne twister rng
    std::uniform_real_distribution<double> distribution(0.0, pertubationStrength);
    for (auto it = Bs.begin(); it != Bs.end(); it++) {
        double abs = 0;
        for (unsigned i = 0; i < it->dimension(1); i++) {
            (*it)(0, i, 0) += distribution(generator);
            abs += std::pow((*it)(0, i, 0).real(), 2);
        }
        abs = std::sqrt(abs);
        for (unsigned i = 0; i < it->dimension(1); i++) {
            (*it)(0, i, 0) /= abs;
        }
    }
}



/**
 * @brief Initialize a MPS using copies of the B tensor
 * @details
 *  A random pertubation up to `pertubationStrength` can be applied to each occupation number in the B tensor,
 *  the tensors are normalized after the pertubations are applied. 
 *  This can help break symmetries, which can make get the simulations stuck.
 */
template<typename T>
MPS<T> makeMPS(const e::Tensor<T, 3>& B, size_t L, double pertubationStrength) {
    std::vector<e::Tensor<T, 3>> Bs(L, B);
    if (pertubationStrength > 0) {
        applyPertubationOnB(Bs, pertubationStrength);
    }
    e::Tensor<T, 1> S(1);
    S(0) = 1;
    std::vector<e::Tensor<T, 1>> Ss(L, S);
    return MPS<T>(std::move(Bs), std::move(Ss));
}


/**
 * @brief Initialize a MPS in spin-up state
 * @details
 *  Return a MPS spin-up MPS with length L and bond dimension 1
 */
template<typename T>
MPS<T> initMPS_spinUp(size_t L, double pertubationStrength=0) {
    assert(L > 0);
    e::Tensor<T, 3> B(1, 2, 1);
    B.setZero();
    B(0, 0, 0) = 1;
    return makeMPS(B, L, pertubationStrength);
}

template<typename T>
MPS<T> initMPS_spinDown(size_t L, double pertubationStrength=0) {
    assert(L > 0);
    e::Tensor<T, 3> B(1, 2, 1);
    B.setZero();
    B(0, 1, 0) = 1;
    return makeMPS(B, L, pertubationStrength);
}

/**
 * @brief Initialize a MPS in spin-right state
 * @details
 *  Return a MPS spin-right MPS with length L and bond dimension 1
 */
template<typename T>
MPS<T> initMPS_spinRight(size_t L, double pertubationStrength=0) {
    assert(L > 0);
    e::Tensor<T, 3> B(1, 2, 1);
    B.setZero();
    T value = std::sqrt(0.5);
    B(0, 0, 0) = value;
    B(0, 1, 0) = value;
    return makeMPS(B, L, pertubationStrength);
}


/**
 * @brief Initialize a MPS with particle number states
 * @details
 *  Return a MPS of dimension `nMax+1`, where each site is occupied by `nParticles`.
 */
template<typename T>
MPS<T> initMPS_nParticles(size_t L, size_t nMax, size_t nParticles, double pertubationStrength=0) {
    assert(L > 0);
    assert(nParticles <= nMax);
    e::Tensor<T, 3> B(1, nMax+1, 1);
    B.setZero();
    B(0, nParticles, 0) = 1;
    return makeMPS(B, L, pertubationStrength);
}
