// Copyright (C) 2016  Jerome Lelong <jerome.lelong@imag.fr>
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include "MultiVariateBasis.h"

using namespace std;

namespace StOpt
{

int ComputeDegree::computeNumberOfElements(const RowMatrixXi &p_tensor, int p_degree) const
{
    int total_elements = 0; // Number of elements with p_total degree smaller than deg
    for (int i = 0; i < p_tensor.rows(); i++)
    {
        int local_degree = count(p_tensor, i);
        total_elements += freedom(p_degree, local_degree) + 1;
    }
    return total_elements;
}

void ComputeDegree::copyPreviousTensor(RowMatrixXi &p_tensor, const RowMatrixXi &p_previousTensor, int p_tensorRow, int p_previousTensorRow) const
{
    for (int j = 0; j < p_previousTensor.cols(); j++)
    {
        p_tensor(p_tensorRow, j) = p_previousTensor(p_previousTensorRow, j);
    }
}

RowMatrixXi ComputeDegree::computeTensor(int p_nVariates, int p_degree) const
{
    if (p_nVariates == 1)
    {
        RowMatrixXi p_tensor(p_degree + 1, p_nVariates);
        for (int i = 0; i < p_degree + 1; i++) p_tensor(i, 0) = i;
        return p_tensor;
    }
    else
    {
        /* Compute the tensor with one variate less */
        RowMatrixXi previousTensor = computeTensor(p_nVariates - 1, p_degree);
        /* Compute the number of rows of p_tensor */
        int nb_elements = computeNumberOfElements(previousTensor, p_degree);
        RowMatrixXi p_tensor(nb_elements, p_nVariates);
        int line = 0;
        /* loop on global degree */
        for (int current_degree = 0 ; current_degree <= p_degree; current_degree++)
        {
            /* loop on  p_partial degree between 0 and current_degree */
            for (int partial_deg = 0 ; partial_deg <= current_degree ; partial_deg++)
            {
                int block_start = 0;
                int block_end = 0;
                /* Determine the first element with global degree = p_degree */
                while (block_start < previousTensor.rows() &&
                        count(previousTensor, block_start) < partial_deg)
                {
                    block_start++;
                }
                block_end = block_start;
                /* Find the last element of the block with global degree = p_degree */
                while (block_end < previousTensor.rows() &&
                        count(previousTensor, block_end) == partial_deg)
                {
                    block_end++;
                }
                block_end--;

                /* loop on the degree of the extra dimension */
                for (int i = freedom(current_degree - 1, partial_deg)  + 1;
                        i <= freedom(current_degree, partial_deg); i++)
                {

                    for (int k = block_start; k <= block_end; k++, line++)
                    {
                        copyPreviousTensor(p_tensor, previousTensor, line, k);
                        p_tensor(line, p_nVariates - 1) = i;
                    }
                }
            }
        }
        return p_tensor;
    }
}

int ComputeDegreeSum::count(const RowMatrixXi &p_tensor, int p_row) const
{
    int deg = 0;
    for (int j = 0; j < p_tensor.cols(); j++)
    {
        deg += p_tensor(p_row, j);
    }
    return deg;
}

int ComputeDegreeSum::freedom(int p_total, int p_partial) const
{
    if (p_partial > p_total) return -1;
    return p_total - p_partial;
}

int ComputeDegreeProd::count(const RowMatrixXi &p_tensor, int p_row) const
{
    int deg = 1;
    bool areAllZero = true;
    for (int j = 0; j < p_tensor.cols(); j++)
    {
        const int power = p_tensor(p_row, j);
        if (areAllZero == true && power > 0) areAllZero = false;
        deg *= std::max(power, 1);

    }
    if (areAllZero == true) return 0;
    return deg;
}

int ComputeDegreeProd::freedom(int p_total, int p_partial) const
{
    if (p_partial > p_total) return -1;
    return p_total / max(p_partial, 1);
}

ComputeDegreeHyperbolic::ComputeDegreeHyperbolic(): m_q(0) { }

ComputeDegreeHyperbolic::ComputeDegreeHyperbolic(double q) : m_q(q) { }

double ComputeDegreeHyperbolic::hyperbolicCount(const RowMatrixXi &p_tensor, int p_row) const
{
    double deg_q = 0.;
    for (int j = 0; j < p_tensor.cols(); j++)
    {
        deg_q += pow(p_tensor(p_row, j), m_q);
    }
    return deg_q;
}

RowMatrixXi ComputeDegreeHyperbolic::computeTensor(int p_nVariates, int p_degree) const
{
    ComputeDegreeSum degreeSum;
    RowMatrixXi Tensor = degreeSum.computeTensor(p_nVariates, p_degree);
    RowMatrixXi TensorHyperbolic(Tensor.rows(), p_nVariates);
    double degree_q = pow(p_degree, m_q);
    int i_sparse = 0;
    for (int i = 0; i < Tensor.rows(); i++)
    {
        if (hyperbolicCount(Tensor, i) > degree_q) continue;
        TensorHyperbolic.row(i_sparse) = Tensor.row(i);
        i_sparse++;
    }
    TensorHyperbolic.conservativeResize(i_sparse, Eigen::NoChange);
    return TensorHyperbolic;
}

}

