// Copyright (C) 2016 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <vector>
#include <memory>
#include <iostream>
#include <Eigen/Dense>
#include "StOpt/regression/LocalSameSizeConstRegression.h"
#include "StOpt/core/grids/InterpolatorSpectral.h"

using namespace std;
using namespace Eigen;

namespace StOpt
{


LocalSameSizeConstRegression::LocalSameSizeConstRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXd &p_lowValues,
        const ArrayXd &p_step,
        const  ArrayXi &p_nbStep) : LocalSameSizeRegression(p_bZeroDate, p_particles, p_lowValues, p_step, p_nbStep)
{
    if (!m_bZeroDate)
    {
        // regression matrix
        m_matReg = ArrayXd::Zero(m_nbMeshTotal);
        for (size_t i = 0; i < m_simAndCell.size(); ++i)
            m_matReg(m_simAndCell[i](1)) += 1;
        // now set diagonal with no particle to 1
        for (int icell = 0; icell < m_nbMeshTotal; ++icell)
            if (m_matReg(icell) <= 0.)
                m_matReg(icell) = 1;
    }
}

LocalSameSizeConstRegression:: LocalSameSizeConstRegression(const   LocalSameSizeConstRegression &p_object): LocalSameSizeRegression(p_object),
    m_matReg(p_object.getMatReg())

{}


void LocalSameSizeConstRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd  &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    if (!m_bZeroDate)
    {

        // update utilitary arrays
        fillInSimCell(m_particles);

        m_matReg = ArrayXd::Zero(m_nbMeshTotal);
        for (size_t i = 0; i < m_simAndCell.size(); ++i)
            m_matReg(m_simAndCell[i](1)) += 1;
        // now set diagonal with no particle to 1
        for (int icell = 0; icell < m_nbMeshTotal; ++icell)
            if (m_matReg(icell) <= 0.)
                m_matReg(icell) = 1;
    }
}

ArrayXd LocalSameSizeConstRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if (!m_bZeroDate)
    {
        ArrayXd ret = ArrayXd::Zero(m_matReg.size());
        // second member
        for (size_t i = 0; i < m_simAndCell.size(); ++i)
            ret(m_simAndCell[i](1)) += p_fToRegress(m_simAndCell[i](0));
        return ret / m_matReg;
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}

ArrayXXd LocalSameSizeConstRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if (!m_bZeroDate)
    {
        ArrayXXd regFunc = ArrayXXd::Zero(p_fToRegress.rows(), m_matReg.size());
        // second member
        for (size_t i = 0; i < m_simAndCell.size(); ++i)
            regFunc.col(m_simAndCell[i](1)) += p_fToRegress.col(m_simAndCell[i](0));

        for (int im = 0; im < m_matReg.size(); ++im)
            regFunc.col(im) /= m_matReg(im);
        return regFunc;
    }
    else
    {
        ArrayXXd retAverage(p_fToRegress.rows(), 1);
        for (int nsm = 0; nsm <  p_fToRegress.rows(); ++nsm)
            retAverage.row(nsm).setConstant(p_fToRegress.row(nsm).mean());
        return retAverage;
    }

}

ArrayXd LocalSameSizeConstRegression::reconstruction(const ArrayXd &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        ArrayXd ret = ArrayXd::Zero(m_particles.cols());
        for (size_t im  = 0 ; im <  m_simAndCell.size(); ++im)
            ret(m_simAndCell[im](0)) = p_basisCoefficients(m_simAndCell[im](1));
        return ret ;
    }
    else
        return ArrayXd::Constant(m_simToCell.size(), p_basisCoefficients(0));
}

ArrayXXd LocalSameSizeConstRegression::reconstructionMultiple(const ArrayXXd &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        ArrayXXd ret = ArrayXXd::Zero(p_basisCoefficients.rows(), m_particles.cols());
        for (size_t im  = 0 ; im <  m_simAndCell.size(); ++im)
            for (int nsm  = 0; nsm < p_basisCoefficients.rows(); ++nsm)
                ret(nsm, m_simAndCell[im](0)) = p_basisCoefficients(nsm, m_simAndCell[im](1));
        return ret ;
    }
    else
    {
        ArrayXXd retValue(p_basisCoefficients.rows(), m_particles.cols());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
            retValue.row(nsm).setConstant(p_basisCoefficients(nsm, 0));
        return retValue ;
    }
}

double LocalSameSizeConstRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    if (!m_bZeroDate)
    {
        if (m_simToCell(p_isim) < 0)
            return 0;
        else
            return p_basisCoefficients(m_simToCell(p_isim)) ;
    }
    else
    {
        return p_basisCoefficients(0);
    }
}


ArrayXd LocalSameSizeConstRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    if (m_bZeroDate)
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());
    // coefficients
    ArrayXd BasisCoefficients = ArrayXd::Zero(m_matReg.size());
    // second member
    for (size_t i = 0; i < m_simAndCell.size(); ++i)
        BasisCoefficients(m_simAndCell[i](1)) += p_fToRegress(m_simAndCell[i](0));
    BasisCoefficients /=  m_matReg;
    ArrayXd  condEspectationValues = ArrayXd::Zero(p_fToRegress.size());
    for (size_t im  = 0 ; im <  m_simAndCell.size(); ++im)
        condEspectationValues(m_simAndCell[im](0)) = BasisCoefficients(m_simAndCell[im](1));
    return condEspectationValues;
}

ArrayXXd LocalSameSizeConstRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    if (m_bZeroDate)
    {
        ArrayXXd ret(p_fToRegress.rows(), p_fToRegress.cols());
        for (int ism = 0; ism < p_fToRegress.rows(); ++ism)
            ret.row(ism).setConstant(p_fToRegress.row(ism).mean());
        return ret;
    }
    // coefficients
    ArrayXXd BasisCoefficients = ArrayXXd::Zero(p_fToRegress.rows(), m_matReg.size());
    // second member
    for (size_t i = 0; i < m_simAndCell.size(); ++i)
        BasisCoefficients.col(m_simAndCell[i](1)) += p_fToRegress.col(m_simAndCell[i](0));

    for (int im = 0; im < m_matReg.size(); ++im)
        BasisCoefficients.col(im) /= m_matReg(im);

    ArrayXXd  condEspectationValues = ArrayXXd::Zero(p_fToRegress.rows(), p_fToRegress.cols());
    for (size_t im  = 0 ; im <  m_simAndCell.size(); ++im)
        condEspectationValues.col(m_simAndCell[im](0)) = BasisCoefficients.col(m_simAndCell[im](1));
    return condEspectationValues;
}

double LocalSameSizeConstRegression::getValue(const ArrayXd &p_coordinates, const ArrayXd &p_coordBasisFunction) const
{
    if (!m_bZeroDate)
    {
        // point location
        int ipos = pointLocation(p_coordinates);
        if (ipos < 0)
            return 0. ;
        else
            return p_coordBasisFunction(ipos);
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double LocalSameSizeConstRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if (!m_bZeroDate)
    {
        // point location
        int ipos = pointLocation(p_coordinates);
        if (ipos < 0)
            return 0. ;
        else
            return  p_interpFuncBasis[ipos]->apply(p_ptOfStock);
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}
}
