#include "mps.hpp"
#include "type.hpp"
#include "util.hpp"
#include <print>


template<typename T>
MPS<T>::MPS(std::vector<BTensor_t>&& Bs, std::vector<STensor_t>&& Ss) {
	this->L = Bs.size();
	assert(this->L > 0);
	assert(Ss.size() == this->L);
	this->Bs = std::move(Bs);
	this->Ss = std::move(Ss);
}

template<typename T>
MPS<T>::BTensor_t MPS<T>::getTheta1(size_t siteIdx) const {
	assert(siteIdx < L);
	auto diag = diagonal(this->S(siteIdx));
	return diag.contract(this->B(siteIdx), e::array<iP, 1>{iP(1, 0)});
}

template<typename T>
e::Tensor<T, 4> MPS<T>::getTheta2(size_t bondIdx) const {
	size_t j = bondIdx + 1;
	assert(j < L);
	auto theta1 = this->getTheta1(bondIdx);
	return theta1.contract(this->B(j), e::array<iP, 1>{iP(2, 0)});
}

template<typename T>
std::vector<unsigned> MPS<T>::getChi() const {
	std::vector<unsigned> chis(this->L-1);
	for (size_t i = 0; i < (L - 1); i++) {
		chis.at(i) = this->B(i).dimension(2);
	}
	return chis;
}
template<typename T>
unsigned MPS<T>::getChiMax() const {
	unsigned chiMax = this->B(0).dimension(2);
	for (size_t i = 1; i < (L - 1); i++) {
		if (this->B(i).dimension(2) > chiMax) {
			chiMax = this->B(i).dimension(2);
		}
	}
	return chiMax;
}
template<typename T>
double MPS<T>::getChiAverage() const {
	double chiAvg = 0;
	for (size_t i = 0; i < (L - 1); i++) {
		chiAvg += this->B(i).dimension(2);
	}
	return chiAvg / (L-1);
}

template<typename T>
T MPS<T>::getSiteExpectationValue(const e::Tensor<T, 2>& op, size_t siteIdx) const {
	auto theta1 = this->getTheta1(siteIdx);
	// optheta: vL, vR, i
	e::Tensor<T, 3> opTheta = theta1.contract(op, e::array<iP, 1>{iP(1, 1)});
	e::Tensor<T, 0> expVal = opTheta.contract(theta1.conjugate(), e::array<iP, 3>{iP(0, 0), iP(2, 1), iP(1, 2)});
	T retVal = expVal(0);
	return retVal;
}

template<typename T>
std::vector<T> MPS<T>::getSiteExpectationValues(const e::Tensor<T, 2>& op) const {
	std::vector<T> results;
	results.reserve(this->L);
	for (size_t i = 0; i < this->L; i++) {
		results.push_back(this->getSiteExpectationValue(op, i));
	}
	return results;
}

template<typename T>
T MPS<T>::getSiteExpectationValueSum(const e::Tensor<T, 2>& op) const {
    dtype expVal = 0;
    for (unsigned i = 0; i < this->getL(); i++) {
        expVal += this->getSiteExpectationValue(op, i);
    }
    return expVal;
}

template<typename T>
T MPS<T>::getBondExpectationValue(const e::Tensor<T, 4>& op, size_t bondIdx) const {
	auto theta2 = this->getTheta2(bondIdx);  // vL i j vR
	// vL [i] [i+1] vR  x  i i+1 [i*] [i+1*]
	e::Tensor<T, 4> opTheta = theta2.contract(op, e::array<iP, 2>{iP{1, 2}, iP{2, 3}});
	// [vL] [vR] [i] [i+1]  x  [vL*] [i*] [i+1*] [vR*]
	e::Tensor<T, 0> expVal = opTheta.contract(theta2.conjugate(), e::array<iP, 4>{iP(0, 0), iP(1, 3), iP(2, 1), iP{3, 2}});
	T retVal = expVal(0);
	return retVal;
}

template<typename T>
T MPS<T>::getEntanglementEntropy(size_t bondIndex) const {
    assert(bondIndex < this->L-1);
    const auto& S = this->S(bondIndex+1);
    STensor_t Ssquared = S * S;
    e::Tensor<dtype, 0> retVal = -(Ssquared * Ssquared.log()).sum();
    return retVal(0);
}


template<typename T>
T MPS<T>::collapse(size_t siteIndex) const {
    // const STensor_t& S = this->S(siteIndex);
    // e::Tensor<T, 0> SS = S.contract(S.conjugate(), e::array<iP, 1>{iP{0,0}});
    // std::println("S[{}] contracted with its conjugate = {}", siteIndex, SS(0));

    auto theta = this->getTheta1(siteIndex);
    // vR vR*
    e::Tensor<dtype, 2> thetaTheta = theta.contract(theta.conjugate(), e::array<iP, 2>{iP{0, 0}, iP{1, 1}});
    for (size_t i = siteIndex+1; i < this->L; i++) {
        // vL vR vL* vR*
		// TODO optimize by not assigning to tensor
        e::Tensor<dtype, 4> BB = this->B(i).contract(this->B(i).conjugate(), e::array<iP, 1>{iP{1, 1}});
        e::Tensor<dtype, 2> thetaTheta2 = thetaTheta.contract(BB, e::array<iP, 2>{iP{0, 0}, iP{1, 2}});
        std::swap(thetaTheta, thetaTheta2);
    }
    e::Tensor<dtype, 1> id(thetaTheta.dimension(0));
    id.setConstant(1);
    auto idDiag = diagonal(id);
    e::Tensor<dtype, 0> retval = thetaTheta.contract(idDiag, e::array<iP, 2>{iP{0, 0}, iP{1, 1}});
	return retval(0);
}

template<typename T>
e::MatrixX<T> tensor2Matrix(const e::Tensor<T, 4>& t) {
    size_t d0 = t.dimension(0);
    size_t d1 = t.dimension(1);
    size_t d2 = t.dimension(2);
    size_t d3 = t.dimension(3);
    e::MatrixX<T> matrix(d0 * d1, d2 * d3);
    for (size_t i0 = 0; i0 < d0; i0++) {
        for (size_t i1 = 0; i1 < d1; i1++) {
            for (size_t i2 = 0; i2 < d2; i2++) {
                for (size_t i3 = 0; i3 < d3; i3++) {
                    // Variant 1 (consistent with what a np.reshape would give)
                    // matrix(i0 * d1 + i1, i2 * d3 + i3) = t(i0, i1, i2, i3);
                    // Variant 2 (consistent with what using Eigen::Map on Tensor<ColumnMajor> produces)
                    matrix(i0 * d1 + i1, i3 * d2 + i2) = t(i0, i1, i2, i3);
                }
            }
        }
    }
    return matrix;
}

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) {
    // Unpack dimensions of the input tensor
    e::Index chivL = theta.dimension(0);
    e::Index dL = theta.dimension(1);
    e::Index dR = theta.dimension(2);
    e::Index chivR = theta.dimension(3);
    // Reshape theta to 2D matrix for the SVD
    using MatX = e::Matrix<dtype, -1, -1>;
    /// todo (DONE): Possible optimization: either use a Map and later adjust the matrix to tensor loops, or globally use RowMajor memory layout?
    // The map approach produces the same Matrix is tensor2Matrix with Variant 2
    MatX theta_mat = e::Map<const MatX>(theta.data(), chivL * dL, dR * chivR);
    // MatX theta_mat = tensor2Matrix(theta);
	return split_truncate_theta(theta_mat, chivL, dL, dR, chivR, chiMax, 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) {
    using MatX = e::Matrix<dtype, -1, -1>;

    // std::cout << "theta:\n" << theta << "theta_mat\n": << theta_mat << std::endl;
    // std::cout << "theta_mat = " << std::endl;
    // std::cout << theta_mat << std::endl;

    // perform SVD: theta = U * diag(S) * V^T
    constexpr int SVDOptions = e::ComputeThinU | e::ComputeThinV;
    e::BDCSVD<MatX, SVDOptions> svd(theta);
    // e::BDCSVD<MatX, SVDOptions> svd(theta_mat);
    const MatX& U = svd.matrixU();
    const e::VectorX<dtype>& S = svd.singularValues();
    const MatX& V = svd.matrixV().transpose();

    // determine the truncation dimension based on chi_max and eps
    int chivC = static_cast<int>(std::min(chiMax, static_cast<size_t>((S.real().array() > std::real(eps)).count())));
    if (chivC < 1) {
        std::println("WARNING: chivC = 0. Setting to 1");
        chivC = 1;
    }
    assert(chivC >= 1);
    assert(chivC <= S.size());

    // NOT NEEDED BECAUSE ITS ALREADY RETURNED IN ORDER
    // sort indices by singular values (descending order) and truncate
    // std::vector<size_t> indices(S.size());
    // std::iota(indices.begin(), indices.end(), 0); // fill indices with [0, 1, 2, ..., Y.size()-1]
    // std::partial_sort(indices.begin(), indices.begin() + chivC, indices.end(),
    //                   [&](size_t a, size_t b) { return S(a) > S(b); });

    // std::println("chivL={}, dL={}, dR={}, chivR={}", chivL, dL, dR, chivR);
    // std::println("theta_mat.shape = {}, {}", theta_mat.rows(), theta_mat.cols());
    // std::println("U.shape = {}, {}", U.rows(), U.cols());
    // std::println("S.shape = {}, {}", S.asDiagonal().rows(), S.asDiagonal().cols());
    // std::println("V.shape = {}, {}", V.rows(), V.cols());
    // std::println("S.shape = {}", S.rows());
    // std::println("chivC={}", chivC);

    
    // TEST
    // MatX theta2 = U * S.asDiagonal() * V;
    // std::cout << "theta2" << std::endl << theta2 << std::endl;
    // for (int i = 0; i < theta2.size(); i++) {
    //     std::print("{} ", theta2.data()[i]);
    // }
    // std::println();

    // std::cout << "U" << std::endl << U << std::endl;
    // std::cout << "U[0, 1]=" << U(0, 1) << std::endl;
    // std::cout << "S" << std::endl << S << std::endl;
    // std::cout << "V" << std::endl << V << std::endl;

    // MatX Xtrunc(dL * chivL, chivC)
    // Truncated versions for debugging
    // MatX Utrunc = U(e::placeholders::all, e::seq(0, chivC));
    // e::VectorXd Strunc = S(e::seq(0, chivC));
    // MatX Vtrunc = V(e::seq(0, chivC), e::placeholders::all);

    // std::cout << "Utrunc" << std::endl << Utrunc << std::endl;
    // std::cout << "Strunc" << std::endl << Strunc << std::endl;
    // std::cout << "Vtrunc" << std::endl << Vtrunc << std::endl;
    // std::println("U.shape = {}, {}", Utrunc.rows(), Utrunc.cols());
    // std::println("S.shape = {}, {}", Strunc.asDiagonal().rows(), Strunc.asDiagonal().cols());
    // std::println("V.shape = {}, {}", Vtrunc.rows(), Vtrunc.cols());

    // MatX theta3 = Utrunc * Strunc.asDiagonal() * Vtrunc;
    // std::cout << "theta3" << std::endl << theta2 << std::endl;
    // for (int i = 0; i < theta3.size(); i++) {
    //     std::print("{} ", theta3.data()[i]);
    // }
    // std::println();

    // reshape X and Z into 3D tensors A and B
    e::Tensor<dtype, 3> A(chivL, dL, chivC);
    e::Tensor<dtype, 3> B(chivC, dR, chivR);
    for (int ivC = 0; ivC < chivC; ivC++) {
        for (int ivL = 0; ivL < chivL; ivL++) {
            for (int idL = 0; idL < dL; idL++) {
                // Variant 1: When the Matrix is constructed with matrix(i0 * d1 + i1) = tensor(i0, i1, i2, i3);
                // dtype value = U(idL + ivL * dL, ivC);
                // std::println("A(ivL={}, idL={}, ivC={}) = U(idL={}+ivL={}*dL={}={}, ivC={}) = {}", 
						// ivL, idL, ivC,    idL, ivL, dL, idL+ivL*dL, ivC, value);

                // Variant 2: When the Matrix is constructed with matrix(i1 * d0 + i0) = tensor(i0, i1, i2, i3);
                dtype value = U(ivL + idL * chivL, ivC);
                // std::println("A(ivL={}, idL={}, ivC={}) = U(ivL={}+idL={}*chivL={}={}, ivC={}) = {}", 
					// ivL, idL, ivC,    ivL, idL, chivL, ivL+idL*chivL, ivC, value
                // );
                A(ivL, idL, ivC) = value;
            }
        }
        for (int ivR = 0; ivR < chivR; ivR++) {
            for (int idR = 0; idR < dR; idR++) {
                // dtype value = V(ivC, ivR + idR * chivR);
                // Variant 1: When the Matrix is constructed with matrix(..., i2 * d3 + i3) = tensor(i0, i1, i2, i3);
                // dtype value = V(ivC, ivR + idR * chivR);
                // std::println("B(ivC={}, idR={}, ivR={}) = V(ivC={}, ivR={}+idR={}*chivR={}={}) = {}", 
                //         ivC, idR, ivR,    ivC, idR, ivR, dR, ivR+idR*chivR, value);
                // Variant 2: When the Matrix is constructed with matrix(..., i3 * d2 + i2) = tensor(i0, i1, i2, i3);
                dtype value = V(ivC, idR + ivR * dR);
                // std::println("B(ivC={}, idR={}, ivR={}) = V(ivC={}, idR={}+ivR={}*dR={}={}) = {}", 
                        // ivC, idR, ivR,    ivC, idR, ivR, dR, idR+ivR*dR, value
                // );
                B(ivC, idR, ivR) = value;
            }
        }
    }
    
    // S: truncate, normalize, and put into a Tensor
    auto norm = S.head(chivC).norm();
    e::Tensor<dtype, 1> Sret(chivC);
    for (int ivC = 0; ivC < chivC; ivC++) {
        Sret(ivC) = S(ivC) / norm;
        // skip normalization for testing if the SVD returns the original matrix after truncation:
        // Sret(ivC) = S(ivC);
    }


    // Verify that A * S * B gives the original matrix back
    // std::cout << "A" << std::endl << A << std::endl;
    // std::println("B.shape = {}, {}, {}", B.dimension(0), B.dimension(1), B.dimension(2));
    // std::cout << "B" << std::endl << B << std::endl;
    // std::println("A.shape = {}, {}, {}", A.dimension(0), A.dimension(1), A.dimension(2));

    // e::Tensor<dtype, 2> Sdiag = diagonal(Sret);
    // std::println("Sdiag.shape = {}, {}", Sdiag.dimension(0), Sdiag.dimension(1));
    // e::Tensor<dtype, 3> theta4 = A.contract(Sdiag, e::array<iP, 1>{iP(2, 0)});
    // e::Tensor<dtype, 4> theta5 = theta4.contract(B, e::array<iP, 1>{iP(2, 0)});

    // std::cout << "theta5" << std::endl << theta5 << std::endl;
    // std::println("theta5.shape = {}, {}, {} {}", theta5.dimension(0), theta5.dimension(1), theta5.dimension(2), theta5.dimension(3));
    
    return {A, Sret, B};
}


// instantiations
template class MPS<dtype>;	
template std::tuple<e::Tensor<dtype, 3>, e::Tensor<dtype, 1>, e::Tensor<dtype, 3>> 
split_truncate_theta<dtype>(const e::Tensor<dtype, 4>&, size_t, dtype);
