// Copyright (C) 2018 EDF
// All Rights Reserved
// This code is published under the GNU Lesser General Public License (GNU LGPL)
#include <vector>
#include <memory>
#include <algorithm>
#include <Eigen/Dense>
#include <iostream>


using namespace std;
using namespace Eigen ;


namespace StOpt
{

//  To ease sorting of points in each dimension
double lesserPair(const pair<double, int > &c1, const pair<double, int> &c2)
{
    return c1.first < c2.first ;
}


///  p_pt           size (2,NnbSimul)
///  p_iSort1       first set of points of size (nbSimul,2)
///  p_iSort2       second set of points of size (nbSimul,2)
///  p_valToAdd     size (P, nbSimul)
///  p_fDominLoc    size (nbFunc, N)
///  merge the two set in first direction for the last dimension
void merge1D1(const  ArrayXXd &p_pt,
              const  ArrayXXi &p_iSort1,
              const  ArrayXXi &p_iSort2,
              const ArrayXXd   &p_valToAdd,
              ArrayXXd   &p_fDominLoc)
{
    int i1 = p_iSort1.rows();
    int i2 =  p_iSort2.rows();
    int iloc = 0 ;
    ArrayXd sumToAdd = ArrayXd::Zero(p_valToAdd.rows());
    for (int i = 0; i < i2 ; ++i)
    {
        int ipt2 = p_iSort2(i, 0);
        int ipt1 = p_iSort1(iloc, 0);
        while (p_pt(0, ipt2) >=  p_pt(0, ipt1))
        {
            sumToAdd += p_valToAdd.col(ipt1);
            iloc += 1;
            if (iloc == i1)
                break;
            ipt1 =  p_iSort1(iloc, 0);
        }
        // add contribution
        p_fDominLoc.col(ipt2) += sumToAdd;
        // last points to treat
        if (iloc == i1)
        {
            for (int j = i + 1; j <  i2; ++j)
            {
                p_fDominLoc.col(p_iSort2(j, 0)) += sumToAdd;
            }
            break;
        }
    }
}



///  p_pt         size (2,nbSimul)
///  p_iSort1     first set of points of size (nbSimul,2)
///  p_iSort2     second set of points of size (nbSimul,2)
///  p_valToAdd   size (P, nbSimul)
///  p_fDominLoc  size (P, N)
///  merge the two set in second direction for the last dimension
void merge1D2(const  ArrayXXd &p_pt,
              const  ArrayXXi &p_iSort1,
              const  ArrayXXi &p_iSort2,
              const ArrayXXd   &p_valToAdd,
              ArrayXXd   &p_fDominLoc)
{
    int i1 = p_iSort1.rows();
    int i2 =  p_iSort2.rows();
    int iloc = i1 - 1 ;
    ArrayXd sumToAdd = ArrayXd::Zero(p_valToAdd.rows());
    for (int i = i2 - 1; i >= 0; --i)
    {
        int ipt2 = p_iSort2(i, 0);
        int ipt1 = p_iSort1(iloc, 0);
        while (p_pt(0, ipt1) >= p_pt(0, ipt2))
        {
            sumToAdd += p_valToAdd.col(ipt1);
            iloc -= 1;
            if (iloc == -1)
                break;
            ipt1 =  p_iSort1(iloc, 0);
        }
        // add contribution
        p_fDominLoc.col(ipt2) += sumToAdd;
        // last points to treat
        if (iloc == -1)
        {
            for (int j = 0; j < i; ++j)
            {
                p_fDominLoc.col(p_iSort2(j, 0)) += sumToAdd;
            }
            break;
        }
    }
}



///  Merge nD procedure between two sets of point to calculate all summations
///  dimension above 2
///  Set A is dominated by set B in the current dimension
///  p_pt          size (d,nbSimul)
///  p_iSort1      first set of points A  size (nbSimul,d)
///  p_iSort2      second set of points B  size (nbSimul,d
///  p_idim        current dimension treated
///  p_valToAdd   vector of size  2^{p_dim}  of arrays of size (P, nbSimul)
///  p_fDomin     vector of size  2^{p_dim}  of arrays of size (P, nbSimul)
void mergeND(const  ArrayXXd &p_pt,
             const  ArrayXXi &p_iSort1,
             const  ArrayXXi &p_iSort2,
             const int &p_idim,
             const std::vector< shared_ptr< ArrayXXd > >   &p_valToAdd,
             std::vector< shared_ptr< ArrayXXd > >   &p_fDomin)
{

    int i1 = p_iSort1.rows();
    int i2 =  p_iSort2.rows();
    // merge the two set to find the median point of the union
    int nbPoints = i1 + i2;
    int nbPtsDiv2 = nbPoints / 2;
    int iPos1 = 0; //position in array1
    int iPos2 = 0 ; // position in array 2
    int iPos = 0;
    int nDimMu = p_idim - 1;
    while ((iPos1 < i1) && (iPos2 < i2) && (iPos < nbPtsDiv2))
    {
        if (p_pt(nDimMu, p_iSort1(iPos1, nDimMu)) < p_pt(nDimMu, p_iSort2(iPos2, nDimMu)))
            iPos1++;
        else
            iPos2++;
        iPos++ ;
    }
    if (iPos1 == i1)
        iPos2 += nbPtsDiv2 - iPos;
    else if (iPos2 == i2)
        iPos1 += nbPtsDiv2 - iPos;
    double xMin = 0. ;
    if (iPos1 > 0)
    {
        if (iPos2 > 0)
            xMin = max(p_pt(nDimMu, p_iSort1(iPos1 - 1, nDimMu)), p_pt(nDimMu, p_iSort2(iPos2 - 1, nDimMu)));
        else
            xMin = p_pt(nDimMu, p_iSort1(iPos1 - 1, nDimMu));
    }
    else
        xMin = p_pt(nDimMu, p_iSort2(iPos2 - 1, nDimMu));

    double xMax = 0;
    if (iPos1 < i1)
    {
        if (iPos2 < i2)
            xMax = min(p_pt(nDimMu, p_iSort1(iPos1, nDimMu)), p_pt(nDimMu, p_iSort2(iPos2, nDimMu)));
        else
            xMax = p_pt(nDimMu, p_iSort1(iPos1, nDimMu));
    }
    else
        xMax =  p_pt(nDimMu, p_iSort2(iPos2, nDimMu));

    // xMed permist to seperate two sets with roughly the same number of particles
    double xMed = 0.5 * (xMin + xMax) ;

    // 4 sets
    ArrayXXi iSort11(iPos1, p_idim)  ; // set A below : A1
    iSort11.col(nDimMu) = p_iSort1.col(nDimMu).head(iPos1);
    ArrayXXi iSort21(iPos2, p_idim) ; // set B below: B1
    iSort21.col(nDimMu) = p_iSort2.col(nDimMu).head(iPos2);
    ArrayXXi iSort12(i1 - iPos1, p_idim); // set A above: A2
    iSort12.col(nDimMu) = p_iSort1.col(nDimMu).tail(i1 - iPos1);
    ArrayXXi iSort22(i2 - iPos2, p_idim) ; // set B above: B2
    iSort22.col(nDimMu) = p_iSort2.col(nDimMu).tail(i2 - iPos2);

    // now keep sorted point on dimension 0
    for (int id = 0 ; id < p_idim - 1; ++id)
    {
        int iloc11 = 0;
        int iloc12 = 0;
        int iloc21 = 0;
        int iloc22 = 0;
        // Set A
        for (int i = 0; i < i1; ++i)
        {
            int ipt = p_iSort1(i, id); // point number
            if (((p_pt(nDimMu, ipt) <= xMed) && (iloc11 < iPos1)) || (iloc12 == iSort12.rows()))
            {
                // under the median
                iSort11(iloc11++, id) = ipt;
            }
            else
            {
                iSort12(iloc12++, id) = ipt;
            }
        }
        // set B
        for (int i = 0; i < i2; ++i)
        {
            int ipt = p_iSort2(i, id); // point number
            if (((p_pt(nDimMu, ipt) <= xMed) && (iloc21 < iPos2)) || (iloc22 == iSort22.rows()))
            {
                // under the median
                iSort21(iloc21++, id) = ipt;
            }
            else
            {
                iSort22(iloc22++, id) = ipt;
            }
        }
    }
    // merge on the two set A1 and B1
    if ((iPos1 > 0) && (iPos2 > 0))
    {
        mergeND(p_pt, iSort11, iSort21, p_idim, p_valToAdd, p_fDomin);
    }

    // merge on teh two set A2 and B2
    if ((iPos1 < i1) && (iPos2 < i2))
    {
        mergeND(p_pt, iSort12, iSort22, p_idim, p_valToAdd, p_fDomin);
    }


    if (p_idim == 2)
    {
        // merge on inferior dimension : we know that point in iSort22 dominate iSort11 in dimension 2 and 3
        if ((iSort11.rows() > 0) && (iSort22.rows() > 0))
        {
            // merge 1D for the  direction  (x_j<x) (y_j<y)  (z_j< z) : 0
            merge1D1(p_pt, iSort11, iSort22,   *p_valToAdd[0],  *p_fDomin[0]) ;

            // merge 1D for the  direction  (x_j>x) (y_j<y)  (z_j< z) : 1
            merge1D2(p_pt, iSort11, iSort22, *p_valToAdd[1], *p_fDomin[1]);

            // merge 1D for the  direction  (x_j<x) (y_j>y)  (z_j> z) : 6
            merge1D1(p_pt, iSort22, iSort11, *p_valToAdd[6], *p_fDomin[6]);

            // merge 1D for the  direction  (x_j>x) (y_j>x)  (z_j> z) : 7
            merge1D2(p_pt,  iSort22, iSort11, *p_valToAdd[7], *p_fDomin[7]);
        }

        if ((iSort12.rows() > 0) && (iSort21.rows() > 0))
        {
            // merge 1D for the  direction  (x_j<x) (y_j>y)  (z_j< z) : 2
            merge1D1(p_pt,  iSort12, iSort21, *p_valToAdd[2],  *p_fDomin[2]) ;

            // merge 1D for the  direction  (x_j>x) (y_j>y)  (z_j< z) : 3
            merge1D2(p_pt,  iSort12, iSort21, *p_valToAdd[3],  *p_fDomin[3]) ;

            // merge 1D for the direction  (x_j<x) (y_j<y) (z_j> z) : 4
            merge1D1(p_pt,  iSort21, iSort12, *p_valToAdd[4],  *p_fDomin[4]) ;

            // merge 1D for the direction  (x_j>x) (y_j<y) (z_j> z) : 5
            merge1D2(p_pt,  iSort21, iSort12, *p_valToAdd[5],  *p_fDomin[5]) ;
        }
    }
    else
    {
        /// merge in dimension below
        int iSizeV = pow(2, p_idim);
        int iSizeVS2 = pow(2, p_idim - 1);
        if ((iSort11.rows() > 0) && (iSort22.rows() > 0))
        {
            // we know a domination of iSort22 on iSort11
            std::vector<shared_ptr< ArrayXXd > > valToAdd1(iSizeV);
            std::vector<shared_ptr< ArrayXXd > > fDomin1(iSizeV);
            for (int id = 0; id < iSizeVS2 ; ++id)
            {
                valToAdd1[id] = p_valToAdd[id];
                fDomin1[id] = p_fDomin[id];
                valToAdd1[id + iSizeVS2] = p_valToAdd[id + 3 * iSizeVS2];
                fDomin1[id + iSizeVS2] = p_fDomin[id + 3 * iSizeVS2];
            }

            mergeND(p_pt, iSort11, iSort22, nDimMu, valToAdd1, fDomin1);

        }
        if ((iSort12.rows() > 0) && (iSort21.rows() > 0))
        {


            // we know that iSort12 dominated in last coordinate and dominates is the previous iSort21
            std::vector<shared_ptr< ArrayXXd > > valToAdd3(iSizeV);
            std::vector<shared_ptr< ArrayXXd > > fDomin3(iSizeV);
            for (int id = 0; id < iSizeVS2 ; ++id)
            {
                valToAdd3[id] = p_valToAdd[id + iSizeV];
                fDomin3[id] = p_fDomin[id + iSizeV];
                valToAdd3[id + iSizeVS2] = p_valToAdd[id + iSizeVS2];
                fDomin3[id + iSizeVS2] = p_fDomin[id + iSizeVS2];
            }
            mergeND(p_pt, iSort21, iSort12, nDimMu, valToAdd3, fDomin3);


        }
    }


}



///  p_pt          size (d,nbSimul)
///  p_iSort      first set of points A  size (nbSimul,d)
///  p_valToAdd   vector of size  2^{d}  of arrays of size (P, nbSimul)
///  p_fDomin     vector of size  2^{d}  of arrays of size (P, nbSimul)
void recursiveCallND(const  ArrayXXd &p_pt,
                     const  ArrayXXi &p_iSort,
                     const std::vector< shared_ptr< ArrayXXd > >   &p_valToAdd,
                     std::vector< shared_ptr< ArrayXXd > > &p_fDomin)
{
    if (p_iSort.cols() == 1)
    {
        for (int is = 1; is <  p_pt.cols(); ++is)
            p_fDomin[0]->col(p_iSort(is, 0)) = p_fDomin[0]->col(p_iSort(is - 1, 0)) + p_valToAdd[0]->col(p_iSort(is - 1, 0)) ;
        for (int is = p_pt.cols() - 2; is >= 0; --is)
            p_fDomin[1]->col(p_iSort(is, 0)) = p_fDomin[1]->col(p_iSort(is + 1, 0)) + p_valToAdd[1]->col(p_iSort(is + 1, 0)) ;
    }
    else if (p_iSort.rows() > 1)
    {
        // split into two part
        int iSize1 = p_iSort.rows() / 2 ;
        int iSize2 = p_iSort.rows() - iSize1  ;
        int nDimM1 = p_iSort.cols() - 1;
        // position valeu of splitting position
        double xMedium = 0.5 * (p_pt(nDimM1, p_iSort(iSize1 - 1, nDimM1)) + p_pt(nDimM1, p_iSort(iSize1, nDimM1)));
        // utilitary for sorted particles
        ArrayXXi iSort1(iSize1, p_iSort.cols());
        ArrayXXi iSort2(iSize2, p_iSort.cols());
        // copy last dimenson
        iSort1.col(nDimM1) = p_iSort.col(nDimM1).head(iSize1);
        iSort2.col(nDimM1) = p_iSort.col(nDimM1).tail(iSize2);
        // two  first dimensions
        for (int id = 0; id < nDimM1; ++id)
        {
            int iLoc1 = 0 ;
            int iLoc2 = 0 ;
            for (int i = 0 ; i < p_iSort.rows() ; ++i)
            {
                int iPoint = p_iSort(i, id) ; // get back point number
                // decide in which set to add the point
                if (((p_pt(nDimM1, iPoint) < xMedium) && (iLoc1 < iSize1)) || (iLoc2 == iSize2))
                    iSort1(iLoc1++, id) = iPoint;
                else
                    iSort2(iLoc2++, id) = iPoint;
            }
        }

        // call on the two set
        recursiveCallND(p_pt, iSort1, p_valToAdd,  p_fDomin);
        recursiveCallND(p_pt, iSort2, p_valToAdd,   p_fDomin);

        // merge nD for the 2 set
        if (p_iSort.cols() > 2)
            mergeND(p_pt, iSort1, iSort2, nDimM1,  p_valToAdd, p_fDomin);
        else
        {
            // 2D merge
            // merge 1D for direction 1
            merge1D1(p_pt, iSort1, iSort2,   *p_valToAdd[0],  *p_fDomin[0]) ;
            // merge 1D for  direction 2
            merge1D2(p_pt, iSort1, iSort2,   *p_valToAdd[1], *p_fDomin[1]);

            // merge 1D for  direction 3
            merge1D1(p_pt, iSort2, iSort1,   *p_valToAdd[2],  *p_fDomin[2]) ;

            // merge 1D for  direction 4
            merge1D2(p_pt, iSort2, iSort1,   *p_valToAdd[3], *p_fDomin[3]) ;

        }

    }

}


/// \brief  Dominance for kernel
/// \param  p_pt        arrays of point coordinates  (d, nbSimul)
/// \param  p_valToAdd  terms to add  (exponential in summation above) : vector of \f$2^d \f  kinds of terms  of size ( P, nbSimul)
/// \param  p_fDomin    result of summation  \f$2^d \f terms of size  ( P, nbSimul)
void nDDominanceKernel(const ArrayXXd &p_pt,
                       const std::vector< shared_ptr< ArrayXXd > >   &p_valToAdd,
                       std::vector< shared_ptr< ArrayXXd >  > &p_fDomin)
{
    int nbSim = p_pt.cols();
    int nDim = p_pt.rows();
    // sort
    ArrayXXi iSort(nbSim, nDim);
    for (int id = 0; id < nDim; ++id)
    {
        vector< std::pair< double, int> >   xSDim(nbSim);
        for (int i = 0; i < nbSim ; ++i)
        {
            xSDim[i] = make_pair(p_pt(id, i), i);
        }
        // sort
        sort(xSDim.begin(), xSDim.end(), lesserPair);
        for (int i = 0; i < nbSim ; ++i)
            iSort(i, id) = xSDim[i].second ;
    }
    for (int ip = 0; ip < pow(2, nDim); ++ip)
        p_fDomin[ip] = make_shared<ArrayXXd>(ArrayXXd::Zero(p_valToAdd[ip]->rows(),  p_valToAdd[ip]->cols()));

    // recursive call with divide and conquer
    recursiveCallND(p_pt, iSort,  p_valToAdd,  p_fDomin);

}

}
