// Copyright (C) 2019 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include "StOpt/regression/LocalKMeansRegression.h"
#include "StOpt/core/utils/constant.h"

using namespace std;
using namespace Eigen;

namespace StOpt
{

/// \brief comparison operator to sort pair
double lesserPairAdapt(const std::pair< double, int > &c1, const std::pair<double, int> &c2)
{
    return c1.first < c2.first ;
}

LocalKMeansRegression::LocalKMeansRegression(const ArrayXi   &p_nbMesh,
        bool p_bRotationAndRecale):
    LocalAdaptCellRegression(p_nbMesh, p_bRotationAndRecale) {}

LocalKMeansRegression::LocalKMeansRegression(const bool &p_bZeroDate,
        const ArrayXXd  &p_particles,
        const ArrayXi   &p_nbMesh,
        bool p_bRotationAndRecale):
    LocalAdaptCellRegression(p_bZeroDate, p_particles, p_nbMesh, p_bRotationAndRecale)
{
    if ((!m_bZeroDate) && (p_nbMesh.size() != 0))
    {
        vector< int > index(m_particles.cols());
        for (size_t i = 0; i < index.size(); ++i)
            index[i] = i;
        ArrayXi coord(m_nbMesh.size());
        m_mesh.resize(p_nbMesh.size(), m_nbMeshTotal);
        m_simulBelongingToCell.resize(m_nbMeshTotal);
        meshing(0, m_particles, index, m_nbMesh, coord, m_mesh, m_simulBelongingToCell);
        for (int i = 0; i < m_nbMeshTotal; ++i)
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
                m_simToCell((*m_simulBelongingToCell[i])[j]) = i;
    }
    else
    {
        m_simToCell.setConstant(0);
        m_nbMeshTotal = 1;
    }
}

LocalKMeansRegression::LocalKMeansRegression(const bool &p_bZeroDate, const  Eigen::ArrayXi &p_nbMesh,
        const   Eigen::Array< std::array< double, 2>, Eigen::Dynamic, Eigen::Dynamic >   &p_mesh, const   ArrayXd &p_meanX,
        const   ArrayXd   &p_etypX, const   MatrixXd   &p_svdMatrix,
        const  bool &p_bRotationAndRecale) :
    LocalAdaptCellRegression(p_bZeroDate, p_nbMesh,   p_mesh, p_meanX,  p_etypX,  p_svdMatrix, p_bRotationAndRecale)
{}

LocalKMeansRegression::LocalKMeansRegression(const LocalKMeansRegression   &p_object) : LocalAdaptCellRegression(p_object)
{}


/// \brief Calculate the local mesh
/// \param  p_particles    particles used for the meshes.
///                        First dimension  : dimension of the problem,
///                        second dimension : the  number of particles
/// \param p_nbMesh       number of meshes in each direction
/// \param p_simToCell    for each simulation, gives its global position in the Cartesian meshing
/// \param p_mesh         describe the mesh generated (first dimension is  the dimension of the problem, the second  dimension is the number of mesh)
///                       For each cell and dimension gives the coordinates min and max of the cell
/// \param p_mesh1D       second representation of the discretization per dimension (conform mesh)
void LocalKMeansRegression::updateSimulations(const bool &p_bZeroDate, const ArrayXXd  &p_particles)
{
    BaseRegression::updateSimulationsBase(p_bZeroDate, p_particles);
    m_simToCell.resize(m_particles.cols());

    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        if (m_particles.rows() != m_nbMesh.size())
        {
            cout << " Dimension nd  of particles of size (nd, nbSimu) is " << m_particles.rows();
            cout << " and   should be equal to the size of the array describing the mesh refinement " << m_nbMesh.transpose() << endl ;
            abort();
        }
        vector< int > index(m_particles.cols());
        for (size_t i = 0; i < index.size(); ++i)
            index[i] = i;
        ArrayXi coord(m_nbMesh.size());
        m_mesh.resize(m_nbMesh.size(), m_nbMeshTotal);
        m_simulBelongingToCell.resize(m_nbMeshTotal);
        meshing(0, m_particles, index, m_nbMesh, coord, m_mesh, m_simulBelongingToCell);
        for (int i = 0; i < m_nbMeshTotal; ++i)
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
                m_simToCell((*m_simulBelongingToCell[i])[j]) = i;
    }
    else
    {
        m_simToCell.setConstant(0);
    }
}

int LocalKMeansRegression::fromCoordTo1D(const ArrayXi &p_nbMesh,  const ArrayXi &p_coord)
{
    int iret =  p_coord(0);
    int idec = p_nbMesh(0);
    for (int i = 1; i < p_nbMesh.size(); ++i)
    {
        iret +=  p_coord(i) * idec;
        idec *= p_nbMesh(i);
    }
    return iret;
}


void  LocalKMeansRegression::listOfAllMesh(const ArrayXi &p_nbMesh,  const ArrayXi &p_coord, const int &p_idim,  std::vector< int > &p_list)
{
    if ((p_idim + 1) == p_nbMesh.size())
    {
        p_list.push_back(fromCoordTo1D(p_nbMesh, p_coord));
    }
    else
    {
        // nest on all mesh
        for (int i = 0;  i < p_nbMesh(p_idim + 1); ++i)
        {
            ArrayXi coordLoc(p_coord);
            coordLoc(p_idim + 1) = i;
            listOfAllMesh(p_nbMesh, coordLoc, p_idim + 1, p_list);
        }
    }
}


void  LocalKMeansRegression::meshing(const int &p_idim, const ArrayXXd  &p_particles, const vector< int > &p_index,
                                     const ArrayXi &p_nbMesh,  const ArrayXi &p_coord,
                                     Array< array< double, 2>, Dynamic, Dynamic >   &p_mesh,
                                     vector<  shared_ptr< std::vector< int> > >    &p_simulBelongingToCell)
{
    assert(static_cast<int>(p_index.size()) >= p_nbMesh(p_idim));
    vector< std::pair< double, int> >   x(p_index.size());
    for (size_t i = 0; i < p_index.size() ; ++i)
    {
        x[i] = make_pair(p_particles(p_idim, p_index[i]), p_index[i]);
    }
    // sort in acendeing order
    sort(x.begin(), x.end(), lesserPairAdapt);

    ArrayXd center(p_nbMesh(p_idim));
    // initialize center
    for (int i = 0; i < m_nbMesh(p_idim); ++i)
    {
        center(i) = x[static_cast<int>((2 * i + 1) * p_index.size() / (2 * p_nbMesh(p_idim)))].first ;
    }
    // store previous
    ArrayXd centerPrev = center;
    ArrayXi num(m_nbMesh(p_idim) + 1); // first number not in mesh
    num(0) = 0;
    num(m_nbMesh(p_idim)) = p_index.size();
    // error
    double error = 10.;
    while (error > tiny)
    {
        // voronoi cell
        int ipos = 0;
        for (int i = 0 ; i < m_nbMesh(p_idim) - 1; ++i)
        {
            double centerAverage = 0.5 * (center(i) + center(i + 1));
            while (x[ipos].first < centerAverage)
                ipos++;
            num(i + 1) = ipos;
        }
        // actualize center
        for (int i = 0; i < m_nbMesh(p_idim); ++i)
        {
            center(i) = 0;
            for (int j = num(i); j < num(i + 1); ++j)
            {
                center(i) += x[j].first;
            }
            center(i) /= (num(i + 1) - num(i));
        }
        error = (center - centerPrev).abs().sum();
        centerPrev = center;
    }
    // next dimension
    ArrayXi coordLoc(p_coord);
    for (int iMesh = 0 ; iMesh < m_nbMesh(p_idim); ++iMesh)
    {
        vector<int > indexMesh;
        indexMesh.reserve(num(iMesh + 1) - num(iMesh));
        for (int j = num(iMesh); j < num(iMesh + 1); ++j)
        {
            indexMesh.push_back(x[j].second);
        }
        coordLoc(p_idim) = iMesh ;
        double coordMin, coordMax;

        if (iMesh == 0)
            coordMin =  - infty;
        else
            coordMin = 0.5 * (center(iMesh - 1) + center(iMesh));
        if (iMesh == m_nbMesh(p_idim) - 1)
            coordMax = infty;
        else
            coordMax = 0.5 * (center(iMesh) + center(iMesh + 1));

        if (p_idim < (p_nbMesh.size() - 1))
        {
            int iret =  coordLoc(0);
            int idec = p_nbMesh(0);
            for (int i = 1; i <= p_idim; ++i)
            {
                iret +=  coordLoc(i) * idec;
                idec *= p_nbMesh(i);
            }
            // nest on all mesh
            std::vector< int > listOfMesh;
            listOfMesh.reserve(p_nbMesh.prod());
            listOfAllMesh(p_nbMesh,  coordLoc, p_idim, listOfMesh);
            // affect all coordinates
            for (size_t im = 0; im < listOfMesh.size(); ++im)
            {
                m_mesh(p_idim, listOfMesh[im])[0] = coordMin;
                m_mesh(p_idim, listOfMesh[im])[1] = coordMax;
            }
            // recursion
            meshing(p_idim + 1, p_particles, indexMesh,  p_nbMesh, coordLoc, p_mesh, p_simulBelongingToCell);
        }
        else
        {
            int icoord = fromCoordTo1D(p_nbMesh, coordLoc);
            p_simulBelongingToCell[icoord] = make_shared< vector< int> >(indexMesh);
            p_mesh(p_idim, icoord)[0] = coordMin;
            p_mesh(p_idim, icoord)[1] = coordMax;
        }
    }
}
int  LocalKMeansRegression::particleToMesh(const Eigen::ArrayXd &p_oneParticle) const
{
    int icell = 0;
    while (icell < m_mesh.size())
    {
        bool bCont = false;
        for (int id = 0; id < m_nbMesh.size(); ++id)
        {
            if ((m_mesh(id, icell)[0] > p_oneParticle(id)) || (m_mesh(id, icell)[1] < p_oneParticle(id)))
            {
                bCont = true;
                break;
            }
        }
        if (!bCont)
            break;
        icell += 1;
    }
    return icell;
}

ArrayXd LocalKMeansRegression::getCoordBasisFunction(const ArrayXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXd basis(m_nbMeshTotal);
        for (int i = 0; i < m_nbMeshTotal; ++i)
        {
            double ret = 0 ;
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
                ret +=  p_fToRegress((*m_simulBelongingToCell[i])[j]);
            ret /= m_simulBelongingToCell[i]->size();
            basis(i) = ret;
        }
        return basis;
    }
    else
    {
        ArrayXd retAverage(1);
        retAverage(0) = p_fToRegress.mean();
        return retAverage;
    }
}


ArrayXXd LocalKMeansRegression::getCoordBasisFunctionMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd basis(p_fToRegress.rows(), m_nbMeshTotal);
        for (int i = 0; i < m_nbMeshTotal; ++i)
        {
            ArrayXd ret = ArrayXd::Zero(p_fToRegress.rows());
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
            {
                for (int ifunc = 0; ifunc < p_fToRegress.rows(); ++ifunc)
                {
                    ret(ifunc) +=  p_fToRegress(ifunc, (*m_simulBelongingToCell[i])[j]);
                }
            }
            ret /= m_simulBelongingToCell[i]->size();
            basis.col(i) = ret;
        }
        return basis;
    }
    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 LocalKMeansRegression::reconstruction(const ArrayXd &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXd ret(m_particles.cols());
        for (int i = 0; i < m_nbMeshTotal; ++i)
        {
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
                ret((*m_simulBelongingToCell[i])[j]) = p_basisCoefficients(i) ;
        }
        return ret;
    }
    else
        return ArrayXd::Constant(m_simToCell.size(), p_basisCoefficients(0));
}

ArrayXXd LocalKMeansRegression::reconstructionMultiple(const ArrayXXd &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        ArrayXXd ret(p_basisCoefficients.rows(), m_particles.cols());
        for (int i = 0; i < m_nbMeshTotal; ++i)
        {
            for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
                for (int ifunc = 0; ifunc < p_basisCoefficients.rows(); ++ifunc)
                    ret(ifunc, (*m_simulBelongingToCell[i])[j]) = p_basisCoefficients(ifunc, i) ;
        }
        return ret;
    }
    else
    {
        ArrayXXd retValue(p_basisCoefficients.rows(), m_simToCell.size());
        for (int nsm = 0; nsm < p_basisCoefficients.rows(); ++nsm)
            retValue.row(nsm).setConstant(p_basisCoefficients(nsm, 0));
        return retValue ;
    }
}


double LocalKMeansRegression::reconstructionASim(const int &p_isim, const ArrayXd   &p_basisCoefficients) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return p_basisCoefficients(m_simToCell(p_isim)) ;
    }
    else
    {
        return p_basisCoefficients(0);
    }
}


ArrayXd LocalKMeansRegression::getAllSimulations(const ArrayXd &p_fToRegress) const
{
    if ((m_bZeroDate) || (m_nbMesh.size() == 0))
        return ArrayXd::Constant(p_fToRegress.size(), p_fToRegress.mean());

    ArrayXd espCond(p_fToRegress.size());
    for (int i = 0; i < m_nbMeshTotal; ++i)
    {
        double ret = 0 ;
        for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
            ret +=  p_fToRegress((*m_simulBelongingToCell[i])[j]);
        ret /= m_simulBelongingToCell[i]->size();
        for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
            espCond((*m_simulBelongingToCell[i])[j]) = ret;
    }

    return espCond;
}

ArrayXXd LocalKMeansRegression::getAllSimulationsMultiple(const ArrayXXd &p_fToRegress) const
{
    if ((m_bZeroDate) || (m_nbMesh.size() == 0))
    {
        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;
    }
    ArrayXXd espCond(p_fToRegress.rows(), p_fToRegress.cols());
    for (int i = 0; i < m_nbMeshTotal; ++i)
    {
        ArrayXd ret = ArrayXd::Zero(p_fToRegress.rows());
        for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
        {
            for (int ifunc = 0; ifunc < p_fToRegress.rows(); ++ifunc)
            {
                ret(ifunc) +=  p_fToRegress(ifunc, (*m_simulBelongingToCell[i])[j]);
            }
        }
        ret /= m_simulBelongingToCell[i]->size();
        for (size_t j = 0; j < m_simulBelongingToCell[i]->size(); ++j)
            espCond.col((*m_simulBelongingToCell[i])[j]) = ret ;
    }
    return espCond;
}

double LocalKMeansRegression::getValue(const ArrayXd &p_coordinates, const ArrayXd &p_coordBasisFunction) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        // rotation
        VectorXd x = m_svdMatrix * ((p_coordinates - m_meanX) / m_etypX).matrix();
        // get back mesh
        return p_coordBasisFunction(particleToMesh(x.array()));
    }
    else
    {
        return p_coordBasisFunction(0);
    }
}

double LocalKMeansRegression::getAValue(const ArrayXd &p_coordinates,  const ArrayXd &p_ptOfStock,
                                        const vector< shared_ptr<InterpolatorSpectral> > &p_interpFuncBasis) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        // rotation
        VectorXd x = m_svdMatrix * ((p_coordinates - m_meanX) / m_etypX).matrix();
        return  p_interpFuncBasis[particleToMesh(x.array())]->apply(p_ptOfStock);
    }
    else
    {
        return p_interpFuncBasis[0]->apply(p_ptOfStock);
    }
}

ArrayXd LocalKMeansRegression::getValuesOneCell(const ArrayXd &, const int &p_cell, const ArrayXXd   &p_foncBasisCoef) const
{
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        return  p_foncBasisCoef.col(p_cell);
    }
    else
        return  p_foncBasisCoef.col(0);
}



ArrayXd LocalKMeansRegression::getCoordBasisFunctionOneCell(const int &p_iCell, const ArrayXd &p_fToRegress) const
{
    ArrayXd retAverage(1);
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        retAverage(0) = 0. ;
        for (size_t i = 0; i <   m_simulBelongingToCell[p_iCell]->size(); ++i)
            retAverage(0) += p_fToRegress((*m_simulBelongingToCell[p_iCell])[i]);
        retAverage(0) /= m_simulBelongingToCell[p_iCell]->size();
    }
    else
    {
        retAverage(0) = p_fToRegress.mean();
    }
    return retAverage;
}


ArrayXXd LocalKMeansRegression::getCoordBasisFunctionMultipleOneCell(const int    &p_iCell, const ArrayXXd &p_fToRegress) const
{
    ArrayXXd retAverage = ArrayXXd::Zero(p_fToRegress.rows(), 1);
    if ((!m_bZeroDate) && (m_nbMesh.size() != 0))
    {
        for (size_t i = 0; i <   m_simulBelongingToCell[p_iCell]->size(); ++i)
            for (int j = 0; j < p_fToRegress.rows(); ++j)
                retAverage(j, i) += p_fToRegress(j, (*m_simulBelongingToCell[p_iCell])[i]);
        retAverage /= m_simulBelongingToCell[p_iCell]->size();
    }
    else
    {
        for (int nsm = 0; nsm <  p_fToRegress.rows(); ++nsm)
            retAverage.row(nsm).setConstant(p_fToRegress.row(nsm).mean());
    }
    return retAverage;
}

}