// This file is part of GQCG-GQCP.
//
// Copyright (C) 2017-2020  the GQCG developers
//
// GQCG-GQCP is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// GQCG-GQCP is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with GQCG-GQCP.  If not, see <http://www.gnu.org/licenses/>.

#pragma once


#include "Mathematical/Representation/SquareMatrix.hpp"
#include "Mathematical/Representation/Tensor.hpp"

#include <iostream>


namespace GQCP {


/**
 *  A square extension of a rank-4 tensor
 *
 *  @tparam _Scalar      the scalar type
 */
template <typename _Scalar>
class SquareRankFourTensor: public Tensor<_Scalar, 4> {
public:
    using Scalar = _Scalar;

    using Base = Tensor<Scalar, 4>;
    using Self = SquareRankFourTensor<Scalar>;


public:
    /*
     *  CONSTRUCTORS
     */

    /**
     *  Default constructor
     */
    SquareRankFourTensor() :
        Base() {}


    /**
     *  Construct a square rank-4 tensor given a dimension
     *
     *  @param dim      the dimension of the rank-4 tensor
     */
    SquareRankFourTensor(const size_t dim) :
        Base(dim, dim, dim, dim) {}


    /**
     *  A basic constructor from a rank-4 GQCP::Tensor that checks if the given tensor is square
     *
     *  @param tensor       the tensor that should be square
     */
    SquareRankFourTensor(const Base& tensor) :
        Base(tensor) {

        // Check if the given tensor is 'square'.
        auto dims = this->dimensions();
        if ((dims[0] != dims[1]) || (dims[1] != dims[2]) || (dims[2] != dims[3])) {
            throw std::invalid_argument("SquareRankFourTensor(Eigen::TensorBase<OtherDerived, AccessLevel>): The given tensor should have equal dimensions in every rank.");
        }
    }


    /**
     *  Constructor from Eigen::Tensor expressions
     *
     *  @tparam ExpDerived      the type of the Eigen expression (normally generated by the compiler)
     *  @tparam AccessLevel     the access level of the TensorBase (normally generated by the compiler)
     *
     *  @param exp              the Eigen::Tensor expression
     */
    template <typename ExpDerived, int AccessLevel>
    SquareRankFourTensor(const Eigen::TensorBase<ExpDerived, AccessLevel>& exp) :
        Self(Base(exp)) {}  // the Base constructor returns the required type for the square-checking constructor


    /*
     *  NAMED CONSTRUCTORS
     */

    /**
     *  Read a square rank-4 tensor from a given file
     *
     *  @param filename     the name of the file to be read in
     *  @param dim          the dimension of the square rank-4 tensor
     */
    static Self FromFile(const std::string& filename, size_t dim) {

        Self result {dim};
        result.setZero();  // make sure that the tensor is initialized to zero values before reading in

        std::ifstream file {filename};
        if (file.is_open()) {
            std::string line;
            while (std::getline(file, line)) {
                std::vector<std::string> splitted_line;  // create a container for the line to be split in

                // Split the line on any whitespace or tabs.
                boost::split(splitted_line, line, boost::is_any_of(" \t"), boost::token_compress_on);

                if (splitted_line.size() != 5) {
                    throw std::runtime_error("SquareRankFourTensor::FromFile(std::string, size_t): Found a line that doesn't contain exactly 5 fields delimited by whitespace.");
                }

                auto i = std::stoi(splitted_line[0]);
                auto j = std::stoi(splitted_line[1]);
                auto k = std::stoi(splitted_line[2]);
                auto l = std::stoi(splitted_line[3]);
                auto value = std::stod(splitted_line[4]);

                result(i, j, k, l) = value;
            }

            file.close();
        } else {
            throw std::runtime_error("SquareRankFourTensor::FromFile(std::string, size_t): Cannot open the given file. Maybe you specified a wrong path?");
        }

        return result;
    }


    /**
     *  Create a zero-initialized `SquareRankFourTensor`.
     * 
     *  @param dim          The dimension of each of the tensor's axes.
     * 
     *  @return A zero `SquareRankFourTensor`.
     */
    static Self Zero(const size_t dim) {

        SquareRankFourTensor<Scalar> T {dim};  // 'Random' initialization happens here.
        T.setZero();
        return T;
    }


    /**
     *  Create a random-initialized `SquareRankFourTensor`, with values uniformly distributed between [-1,1].
     * 
     *  @param dim          The dimension of each of the tensor's axes.
     * 
     *  @return A random `SquareRankFourTensor`.
     */
    static Self Random(const size_t dim) {

        SquareRankFourTensor<Scalar> T {dim};
        T.setRandom();  // Uniformly distributed between [0, 1].

        // Move the distribution from [0, 1] -> [-1, 1].
        for (size_t i = 0; i < dim; i++) {
            for (size_t j = 0; j < dim; j++) {
                for (size_t k = 0; k < dim; k++) {
                    for (size_t l = 0; l < dim; l++) {
                        T(i, j, k, l) = 2 * T(i, j, k, l) - 1;  // Scale from [0, 1] -> [0, 2] -> [-1, 1].
                    }
                }
            }
        }
        return T;
    }


    /*
     *  OPERATORS
     */

    /**
     *  Assignment operator from Eigen::Tensor expressions
     *
     *  @tparam ExpDerived      the type of the Eigen expression (normally generated by the compiler)
     *  @tparam AccessLevel     the access level of the TensorBase (normally generated by the compiler)
     *
     *  @param exp              the Eigen::Tensor expression
     */
    template <typename ExpDerived, int AccessLevel>
    Self& operator=(const Eigen::TensorBase<ExpDerived, AccessLevel>& exp) {
        this->Base::operator=(exp);
        return (*this);
    }


    /*
     *  PUBLIC METHODS
     */

    using Base::dimension;

    /**
     *  @return the dimension of this square rank-four tensor
     */
    size_t dimension() const { return this->dimension(0); }  // all tensor dimensions are equal because of the constructor

    /**
     *  @return the pair-wise reduction of this square rank-4 tensor, i.e. the tensor analog of a strict "lower triangle" as a matrix in column major form
     *
     *  The matrix indices (m,n) come from the tensor indices (i,j,k,l) and are such that:
     *      - m is compounded in a column major way from i and j, with the restriction i>j
     *      - n is compounded in a column major way from k and l, with the restriction k>l
     */
    SquareMatrix<double> pairWiseStrictReduced() const {

        // Initialize the resulting matrix
        const auto K = this->dimension();
        SquareMatrix<double> M = SquareMatrix<double>::Zero(K * (K - 1) / 2);

        // Calculate the compound indices and bring the elements from the tensor over into the matrix
        size_t row_index = 0;
        for (size_t j = 0; j < K; j++) {          // "column major" ordering for row_index<-i,j so we do j first, then i
            for (size_t i = j + 1; i < K; i++) {  // in column major indices, columns are contiguous, so the first of two indices changes more rapidly
                // require i > j for "lower triangle"

                size_t column_index = 0;
                for (size_t l = 0; l < K; l++) {          // "column major" ordering for column_index<-k,l so we do l first, then k
                    for (size_t k = l + 1; k < K; k++) {  // in column major indices, columns are contiguous, so the first of two indices changes more rapidly
                        // require l > k for "lower triangle"

                        M(row_index, column_index) = this->operator()(i, j, k, l);

                        column_index++;
                    }
                }

                row_index++;
            }
        }

        return M;
    }
};

}  // namespace GQCP
