#include "bose_hubbard.hpp"
#include <eigen3/unsupported/Eigen/KroneckerProduct>


template<typename T>
LadderOperators<T>::LadderOperators(unsigned dim) : identity(dim, dim), aDagger(dim, dim), a(dim, dim), n(dim, dim) {
    identity.setZero();
    aDagger.setZero();
    a.setZero();
    n.setZero();
    identity(0, 0) = 1;

    for (unsigned i = 1; i < dim; i++) {
        identity(i, i) = 1;
        n(i, i) = i;
        // TODO what if T = std::complex<float>?
        double sqrtN = std::sqrt(static_cast<double>(i));
        a(i-1, i) = sqrtN;
        aDagger(i, i-1) = sqrtN;
    }
}


template<typename T>
void BoseHubbardModel1D<T>::initH_Bonds() {
    T UHalf = this->U * static_cast<T>(0.5);
    T UQuarter = this->U * static_cast<T>(0.25);
    T muHalf = this->mu * static_cast<T>(0.5);

    /// @todo TODO figure out which is right. Doesnt t need to be halfed at the edges and not the middle?
    e::MatrixX<T> N_NminusI = this->lo.getN() * (this->lo.getN() - this->lo.getIdentity());

    e::MatrixX<T> twoParticlePart =
        - this->t * e::kroneckerProduct(this->lo.getADagger(), this->lo.getA())
        - this->t * e::kroneckerProduct(this->lo.getA(), this->lo.getADagger()
    );
    auto left =   std::make_shared<e::MatrixX<T>>(
        twoParticlePart
        + e::kroneckerProduct(N_NminusI, this->lo.getIdentity()) * UHalf
        + e::kroneckerProduct(this->lo.getIdentity(), N_NminusI) * UQuarter
        - e::kroneckerProduct(this->lo.getN(), this->lo.getIdentity()) * (this->mu)
        - e::kroneckerProduct(this->lo.getIdentity(), this->lo.getN()) * (muHalf)
    );
    auto right =  std::make_shared<e::MatrixX<T>>(
        twoParticlePart
        + e::kroneckerProduct(N_NminusI, this->lo.getIdentity()) * UQuarter
        + e::kroneckerProduct(this->lo.getIdentity(), N_NminusI) * UHalf
        - e::kroneckerProduct(this->lo.getN(), this->lo.getIdentity()) * (muHalf)
        - e::kroneckerProduct(this->lo.getIdentity(), this->lo.getN()) * (this->mu)
    );
    auto middle = std::make_shared<e::MatrixX<T>>(
        twoParticlePart
        + e::kroneckerProduct(N_NminusI, this->lo.getIdentity()) * UQuarter
        + e::kroneckerProduct(this->lo.getIdentity(), N_NminusI) * UQuarter
        - e::kroneckerProduct(this->lo.getN(), this->lo.getIdentity()) * (muHalf)
        - e::kroneckerProduct(this->lo.getIdentity(), this->lo.getN()) * (muHalf)
    );
    // The matrices are N*N x N*N -> reshape to N x N x N x N tensors
    unsigned N = this->localDim;
    auto leftT = std::make_shared<e::Tensor<T, 4>>(N, N, N, N);
    auto middleT = std::make_shared<e::Tensor<T, 4>>(N, N, N, N);
    auto rightT = std::make_shared<e::Tensor<T, 4>>(N, N, N, N);
    for (unsigned i1 = 0; i1 < N; i1++) {
        for (unsigned i2 = 0; i2 < N; i2++) {
            for (unsigned i3 = 0; i3 < N; i3++) {
                for (unsigned i4 = 0; i4 < N; i4++) {
                    (*leftT)(i1, i2, i3, i4) = (*left)(i1 * N + i2, i3 * N + i4);
                    (*middleT)(i1, i2, i3, i4) = (*middle)(i1 * N + i2, i3 * N+ i4);
                    (*rightT)(i1, i2, i3, i4) = (*right)(i1 * N + i2, i3 * N + i4);
                }
            }
        }
    }
    this->H_Bonds.reserve(this->L-1);
    this->H_Bonds.push_back(leftT);
    for (unsigned i = 0; i < this->L - 3; i++) {  // there are L-3 middle bonds
        this->H_Bonds.push_back(middleT);
    }
    this->H_Bonds.push_back(rightT);

    // TODO: merge or remove tensor version
    this->H_BondsMatrices.reserve(this->L-1);
    this->H_BondsMatrices.push_back(left);
    for (unsigned i = 0; i < this->L - 3; i++) {  // there are L-3 middle bonds
        this->H_BondsMatrices.push_back(middle);
    }
    this->H_BondsMatrices.push_back(right);
}


template<typename T>
void BoseHubbardModel1D<T>::initH_MPO() {
    e::MatrixX<T> zero(this->localDim, this->localDim);
    zero.setZero();
    
    auto W = std::make_shared<e::Tensor<T, 4>>(std::move(constructMPO<T>({
                {this->lo.getIdentity(), this->lo.getADagger(), this->lo.getA(), this->U / 2. * (this->lo.getN() * this->lo.getN() - this->lo.getN()) - this->mu * this->lo.getN()},
                {zero, zero, zero, (-this->t * this->lo.getA())},
                {zero, zero, zero, (-this->t * this->lo.getADagger())},
                {zero, zero, zero, this->lo.getIdentity()}
            })));
    // same for all sites
    this->H_MPO = std::vector<std::shared_ptr<e::Tensor<T, 4>>>(this->L, W);
}

template class LadderOperators<dtype>;
template class BoseHubbardModel1D<dtype>;
