#include "transverse_field_ising.hpp"

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

template<typename T> 
PauliMatrices<T> TransverseFieldIsingModel1D<T>::pm;

template<typename T>
void TransverseFieldIsingModel1D<T>::initH_Bonds() {
    T BHalf = this->B * static_cast<T>(0.5);
    ///@todo TODO remove this after testing!
    // T BHalf = this->B;
    // TODO: improve code or directly initialize the tensors
    auto left =   std::make_shared<e::MatrixX<T>>(
            e::kroneckerProduct(pm.getSigmaZ(), pm.getSigmaZ()) * (-J)
            - this->B * e::kroneckerProduct(pm.getSigmaX(), pm.getIdentity())
            - BHalf * e::kroneckerProduct(pm.getIdentity(), pm.getSigmaX())
    );
    auto right =  std::make_shared<e::MatrixX<T>>(
            e::kroneckerProduct(pm.getSigmaZ(), pm.getSigmaZ()) * (-J)
            - BHalf * e::kroneckerProduct(pm.getSigmaX(), pm.getIdentity())
            - this->B * e::kroneckerProduct(pm.getIdentity(), pm.getSigmaX())
    );
    auto middle = std::make_shared<e::MatrixX<T>>(
            e::kroneckerProduct(pm.getSigmaZ(), pm.getSigmaZ()) * (-J)
            - BHalf * e::kroneckerProduct(pm.getSigmaX(), pm.getIdentity())
            - BHalf * e::kroneckerProduct(pm.getIdentity(), pm.getSigmaX())
    );
    auto leftT = std::make_shared<e::Tensor<T, 4>>(2, 2, 2, 2);
    auto middleT = std::make_shared<e::Tensor<T, 4>>(2, 2, 2, 2);
    auto rightT = std::make_shared<e::Tensor<T, 4>>(2, 2, 2, 2);
    for (unsigned i1 = 0; i1 < 2; i1++) {
        for (unsigned i2 = 0; i2 < 2; i2++) {
            for (unsigned i3 = 0; i3 < 2; i3++) {
                for (unsigned i4 = 0; i4 < 2; i4++) {
                    (*leftT)(i1, i2, i3, i4) = (*left)(i1 * 2 + i2, i3 * 2 + i4);
                    (*middleT)(i1, i2, i3, i4) = (*middle)(i1 * 2 + i2, i3 * 2+ i4);
                    (*rightT)(i1, i2, i3, i4) = (*right)(i1 * 2 + i2, i3 * 2 + i4);
                }
            }
        }
    }
    // *leftT = Bond_Matrix2Tensor(*left, 2);
    // *middleT = Bond_Matrix2Tensor(*middle, 2);
    // *rightT = Bond_Matrix2Tensor(*right, 2);
    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 TransverseFieldIsingModel1D<T>::initH_MPO() {
    e::Matrix2<T> zero(2, 2);
    zero.setZero();
    auto W = std::make_shared<e::Tensor<T, 4>>(std::move(constructMPO<T>({
                {this->pm.getIdentity(), this->pm.getSigmaZ(), (-this->B) * this->pm.getSigmaX()},
                {zero, zero, (-this->J * this->pm.getSigmaZ())},
                {zero, zero, this->pm.getIdentity()}
            })));
    // same for all sites
    this->H_MPO = std::vector<std::shared_ptr<e::Tensor<T, 4>>>(this->L, W);
}

template class TransverseFieldIsingModel1D<dtype>;
