#include "model.hpp"
#include "type.hpp"
#include "util.hpp"

template<typename T>
T Model1D<T>::energy(const MPS<T>& psi) const {
    assert(psi.getL() == this->L);
    T E = 0;
    for (unsigned i = 0; i < this->L-1; i++) {
        E += psi.getBondExpectationValue(*(this->H_Bonds.at(i)), i);
    }
    return E;
}


template<typename T>
e::Tensor<T, 4> constructMPO(const std::initializer_list<std::initializer_list<e::MatrixX<T>>>& matrices) {
    size_t rows = matrices.size();
    assert(rows > 0);
    size_t cols = matrices.begin()->size();
    assert(rows > 1);
    size_t localDim = matrices.begin()->begin()->rows();
    // TODO: check every matrix? 
    assert(localDim == matrices.begin()->begin()->cols());
    e::Tensor<T, 4> mpo(rows, cols, localDim, localDim);
    // for (const auto& [rowIdx, row] : std::ranges::views::enumerate(matrices)) {
    size_t rowIdx = 0;
    for (const auto& row : matrices) {
        assert(row.size() == cols);
        size_t colIdx = 0;
        // for (const auto& [colIdx, mat] : std::ranges::views::enumerate(row)) {
        for (const auto& mat : row) {
            for (size_t i = 0; i < localDim; i++) {
                for (size_t j = 0; j < localDim; j++) {
                    mpo(rowIdx, colIdx, i, j) = mat(i, j);
                }
            }
            colIdx++;
        }
        rowIdx++;
    }
    return mpo;
}


template<typename T>
e::MatrixX<T> MPO_Tensor2Matrix(const e::Tensor<T, 4>& mpo) {
    size_t dim = mpo.dimension(0);
    assert(mpo.dimension(1) == dim);
    size_t localDim = mpo.dimension(2);
    assert(mpo.dimension(3) == localDim);
    e::MatrixX<T> matrix(dim * localDim, dim * localDim);
    for (size_t rowIdx = 0; rowIdx < dim; rowIdx++) {
        for (size_t colIdx = 0; colIdx < dim; colIdx++) {
            for (size_t i = 0; i < localDim; i++) {
                for (size_t j = 0; j < localDim; j++) {
                    matrix(rowIdx * localDim + i, colIdx * localDim + j) = mpo(rowIdx, colIdx, i, j);
                }
            }
        }
    }
    return matrix;
}


template<typename T>
e::Tensor<T, 4> bondMatrix2Tensor(const e::MatrixX<T>& bondMatrix, unsigned localDim) {
    e::Index matDim = bondMatrix.rows();
    assert(bondMatrix.cols() == matDim);
    assert(localDim > 0);
    assert(matDim % localDim == 0);
    e::Index tensorDim = matDim / localDim;
    e::Tensor<T, 4> bondTensor(tensorDim, tensorDim, localDim, localDim);
    for (unsigned rowIdx = 0; rowIdx < tensorDim; rowIdx++) {
        for (unsigned colIdx = 0; colIdx < tensorDim; colIdx++) {
            for (unsigned i = 0; i < localDim; i++) {
                for (unsigned j = 0; j < localDim; j++) {
                    bondTensor(rowIdx, i, colIdx, j) = bondMatrix(rowIdx * localDim + i, colIdx * localDim + j);
                }
            }
        }
    }
    return bondTensor;
}


template e::Tensor<dtype, 4> constructMPO<dtype>(const std::initializer_list<std::initializer_list<e::MatrixX<dtype>>>& matrices);
template e::MatrixX<dtype> MPO_Tensor2Matrix<dtype>(const e::Tensor<dtype, 4>& mpo);
template e::Tensor<dtype, 4> bondMatrix2Tensor<dtype>(const e::MatrixX<dtype>&, unsigned);
template class Model1D<dtype>;
