// Copyright (C) 2016,2021  EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <vector>
#include <array>
#include <memory>
#include <iostream>
#include <algorithm>
#include <Eigen/Dense>
#include "StOpt/regression/meshToCoordinates.h"
#include "StOpt/core/utils/constant.h"

using namespace std;
using namespace Eigen ;

namespace StOpt
{

// Local function that gives the minimum number of particle by cell
int nbMinPartByCell(const ArrayXi &p_simToCell, const int &p_nbcell)
{
    ArrayXi nbPartPerCell = ArrayXi::Zero(p_nbcell);
    for (int is = 0; is <  p_simToCell.size(); ++is)
    {
        nbPartPerCell(p_simToCell(is)) += 1;
    }
    return nbPartPerCell.minCoeff() ;
}

void meshCalculationLocalRegression(const ArrayXXd &p_particles, const ArrayXi   &p_nbMesh,
                                    ArrayXi &p_simToCell, Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
                                    vector< shared_ptr< ArrayXd > > &p_mesh1D)
{
    int nDim =  p_nbMesh.size();
    assert(p_particles.rows() == nDim);

    // number of meshes
    int nbMeshGlob  =  p_nbMesh.prod();
    p_mesh.resize(nDim, nbMeshGlob);
    p_mesh1D.resize(p_nbMesh.size());
    for (int id = 0; id < nDim; ++id)
        p_mesh1D[id] = make_shared< ArrayXd >(p_nbMesh(id) + 1);
    // number of simulations
    int nbSimul = p_particles.cols();
    // utilitary for position in meshes
    int idecCell = 1 ;
    p_simToCell.setConstant(0) ;
    // calculate the meshing (cost linear in simulation number)
    for (int id = 0 ; id < nDim ; ++id)
    {
        vector<double> partDim(nbSimul);
        for (int ip = 0; ip < nbSimul; ++ip)
            partDim[ip] = p_particles(id, ip);
        vector<double>::iterator startD = partDim.begin();
        vector<double>::iterator endD = partDim.end();
        // number of particles per mesh
        int nppm = nbSimul / p_nbMesh(id);
        // partial sort only
        nth_element(startD, startD, endD);
        (*p_mesh1D[id])(0) = partDim[0];
        nth_element(startD + 1, startD + nppm - 1, endD);
        (*p_mesh1D[id])(1) = partDim[nppm - 1];
        for (int j = 1 ; j < p_nbMesh(id) - 1 ; ++j)
        {
            assert((*p_mesh1D[id])(j) >= (*p_mesh1D[id])(j - 1));
            nth_element(startD + j * nppm, startD + (j + 1)*nppm - 1, endD);
            (*p_mesh1D[id])(j + 1) = partDim[(j + 1) * nppm - 1];
        }
        nth_element(startD + (p_nbMesh(id) - 1)*nppm, endD - 1, endD);
        (*p_mesh1D[id])(p_nbMesh(id)) = partDim[nbSimul - 1];
        assert((*p_mesh1D[id])(p_nbMesh(id))  >= (*p_mesh1D[id])(p_nbMesh(id) - 1));
        for (int is = 0 ; is < nbSimul ; ++is)
        {
            int im = 0 ;
            while (p_particles(id, is) > (*p_mesh1D[id])(im + 1)) im++;
            p_simToCell(is) += idecCell * im;
        }
        // offset
        idecCell *= p_nbMesh(id);
    }
    // Utilitary for coordinates
    VectorXi  coord(nDim) ;
    // utilitary
    int nDivInit = ((nDim > 1) ? p_nbMesh.head(nDim - 1).prod() : 1) ;
    // for each mesh  calculate its borders
    for (int im = 0 ; im < nbMeshGlob ; ++im)
    {
        ArrayXi coord = meshToCoordinates(p_nbMesh, nDivInit, im);
        for (int j = 0 ; j < nDim ; ++j)
        {
            // minimal value
            p_mesh(j, im)[0] = (*p_mesh1D[j])(coord(j)) ;
            // maximal value
            p_mesh(j, im)[1] = (*p_mesh1D[j])(coord(j) + 1) ;
        }
    }
    //assert(nbMinPartByCell(p_simToCell, nbMeshGlob) > nDim);
}


void meshCalculationLocalRegressionDiscrLastDim(const ArrayXXd &p_particles, const ArrayXi   &p_nbMesh,
        ArrayXi &p_simToCell, Array< array< double, 2>, Dynamic, Dynamic > &p_mesh,
        vector< shared_ptr< ArrayXd > > &p_mesh1D)
{
    int nDim =  p_nbMesh.size();
    assert(p_particles.rows() == nDim);

    // number of meshes
    int nbMeshGlob  =  p_nbMesh.prod();
    p_mesh.resize(nDim, nbMeshGlob);
    p_mesh1D.resize(p_nbMesh.size());
    for (int id = 0; id < nDim; ++id)
        p_mesh1D[id] = make_shared< ArrayXd >(p_nbMesh(id) + 1);
    // number of simulations
    int nbSimul = p_particles.cols();
    // utilitary for position in meshes
    int idecCell = 1 ;
    p_simToCell.setConstant(0) ;
    // calculate the meshing (cost linear in simulation number)
    // do it only for the first nDim-1 dimension
    for (int id = 0 ; id < nDim - 1 ; ++id)
    {
        vector<double> partDim(nbSimul);
        for (int ip = 0; ip < nbSimul; ++ip)
            partDim[ip] = p_particles(id, ip);
        vector<double>::iterator startD = partDim.begin();
        vector<double>::iterator endD = partDim.end();
        // number of particles per mesh
        int nppm = nbSimul / p_nbMesh(id);
        // partial sort only
        nth_element(startD, startD, endD);
        (*p_mesh1D[id])(0) = partDim[0];
        nth_element(startD + 1, startD + nppm - 1, endD);
        (*p_mesh1D[id])(1) = partDim[nppm - 1];
        for (int j = 1 ; j < p_nbMesh(id) - 1 ; ++j)
        {
            assert((*p_mesh1D[id])(j) >= (*p_mesh1D[id])(j - 1));
            nth_element(startD + j * nppm, startD + (j + 1)*nppm - 1, endD);
            (*p_mesh1D[id])(j + 1) = partDim[(j + 1) * nppm - 1];
        }
        nth_element(startD + (p_nbMesh(id) - 1)*nppm, endD - 1, endD);
        (*p_mesh1D[id])(p_nbMesh(id)) = partDim[nbSimul - 1];
        assert((*p_mesh1D[id])(p_nbMesh(id))  >= (*p_mesh1D[id])(p_nbMesh(id) - 1));
        for (int is = 0 ; is < nbSimul ; ++is)
        {
            int im = 0 ;
            while (p_particles(id, is) > (*p_mesh1D[id])(im + 1)) im++;
            p_simToCell(is) += idecCell * im;
        }
        // offset
        idecCell *= p_nbMesh(id);
    }
    // nb cell last dim
    int nbCellLDim = p_nbMesh(nDim - 1);
    // treat last dimension where values are discrete : there are gathered
    vector<std::pair<double, int> > partDim(nbSimul);
    for (int ip = 0; ip < nbSimul; ++ip)
        partDim[ip] = make_pair(p_particles(nDim - 1, ip), ip);
    // heres sort
    sort(partDim.begin(), partDim.end());
    std::vector< std::shared_ptr< std::vector<int> >  > isimByDiscr;
    isimByDiscr.reserve(10 * nbCellLDim);
    std::shared_ptr< std::vector<int> > ptLoc = make_shared< std::vector<int>>();
    ptLoc->reserve(nbSimul);
    ptLoc->push_back(partDim[0].second);
    for (int ip = 1; ip < nbSimul; ++ip)
    {
        if (std::abs(partDim[ip].first - partDim[ip - 1].first) < tiny)
        {
            ptLoc->push_back(partDim[ip].second);
        }
        else
        {
            isimByDiscr.push_back(ptLoc);
            ptLoc = make_shared< std::vector<int>>();
            ptLoc->reserve(nbSimul);
            ptLoc->push_back(partDim[ip].second);
        }
    }
    if (ptLoc->size() > 0)
    {
        isimByDiscr.push_back(ptLoc);
    }
    int nbDiscrete = isimByDiscr.size();
    int nbSimuOptByCell = static_cast<int>(0.8 * nbSimul / nbCellLDim);
    int nbSimCur = 1;
    size_t iPos = 0 ; // in isimByDiscr
    (*p_mesh1D[nDim - 1])(0) = p_particles(nDim - 1, (*isimByDiscr[0])[0])  - small;
    for (int icell = 0; icell < nbCellLDim; ++icell)
    {
        while (nbSimCur < nbSimuOptByCell)
        {
            for (size_t ip = 0; ip < isimByDiscr[iPos]->size(); ++ip)
            {
                p_simToCell((*isimByDiscr[iPos])[ip]) += idecCell * icell;
            }
            nbSimCur += isimByDiscr[iPos]->size();
            iPos += 1;
            if (iPos == isimByDiscr.size())
                break;
        }
        nbSimCur = 0;
        if ((iPos < isimByDiscr.size()) && (icell < nbCellLDim - 1))
            (*p_mesh1D[nDim - 1])(icell + 1) =  0.5 * (p_particles(nDim - 1, (*isimByDiscr[iPos])[0]) + p_particles(nDim - 1, (*isimByDiscr[iPos - 1])[0])) ;
        else
            break;
    }
    (*p_mesh1D[nDim - 1])(nbCellLDim) = p_particles(nDim - 1, (*isimByDiscr[nbDiscrete - 1])[0]) + small;

    // Utilitary for coordinates
    VectorXi  coord(nDim) ;
    // utilitary
    int nDivInit = ((nDim > 1) ? p_nbMesh.head(nDim - 1).prod() : 1) ;
    // for each mesh  calculate its borders
    for (int im = 0 ; im < nbMeshGlob ; ++im)
    {
        ArrayXi coord = meshToCoordinates(p_nbMesh, nDivInit, im);
        for (int j = 0 ; j < nDim ; ++j)
        {
            // minimal value
            p_mesh(j, im)[0] = (*p_mesh1D[j])(coord(j)) ;
            // maximal value
            p_mesh(j, im)[1] = (*p_mesh1D[j])(coord(j) + 1) ;
        }
    }
    //assert(nbMinPartByCell(p_simToCell, nbMeshGlob) >= nDim - 1);
}
}
