#include <stdlib.h>
#include <stddef.h>
#include "GeoMorph.h"
#include "FloatRaster.h"
#include "FlowDirGrid.h"
#include "Topology.h"
#include "Mesh.h"
#include "MeshTree.h"
#include "esriheader.h"
#include "Erreur.h"
#include <set>
#include <map>
#include <algorithm>

using namespace std;

int GetCatchmentSum_ ( FlowDirGrid * F , FloatRaster * R , int * _outlet, set<int> & UpNodes, int ** _BB , float * _Sum , set<int> & ixUpNodes , int * _nr , int * _nc );
int GetUpstreamNodes_ ( FlowDirGrid * F , int * _outlet, set<int> & UpNodes, set<int> & ixUpNodes );

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int MakeMeshTree ( FlowDirGrid * F , FloatRaster * CW , FloatRaster * PK , int nnodes , int * NODES , int * CNT , MeshTree ** _MTree , double * _meanweight )
{

    int i,j,ichild,nr,nc,nchildren ;
    double SSE , meanweight , weight ;
    float f;
    set<int> ixUpNodes , UpNodes ;
    set<int>::iterator isi ;
    map<float,MeshTree *>MapChildren;
    map<float,MeshTree *>::reverse_iterator MapChildrenIterator;

    (*_MTree) = (MeshTree *) malloc ( nnodes * sizeof ( MeshTree ) );

    UpNodes.clear();

    for(i=0;i<nnodes;i++)
    {
        UpNodes.insert(NODES[i]) ;
        (*_MTree)[i].RootCell = NODES[i] ;
        (*_MTree)[i].Parent = NULL ;
        (*_MTree)[i].SSE = 0 ;
        (*_MTree)[i].Weight = 0 ;
        (*_MTree)[i].cnt = CNT[i] ;
        (*_MTree)[i].PositionRange = NULL ;
        (*_MTree)[i].CWeight = NULL ;
        (*_MTree)[i].RangeLength = 0 ;
        (*_MTree)[i].CurrentPosition = -1 ;

    }

    // Construction de la topologie

    meanweight = 0 ;

    for(i=0;i<nnodes;i++)
    {
        ixUpNodes.clear();

        if ( GetUpstreamNodes_ ( F , NODES + i , UpNodes, ixUpNodes ) )
        {
            SetErreur("MakeMeshTree : echec de calcul de la somme de bassin au noeud # %d\n", NODES[i] );
            return(-1);
        }

        nchildren = ixUpNodes.size();
        (*_MTree)[i].nchildren = nchildren ;
        (*_MTree)[i].Child = (MeshTree **) malloc ( nchildren * sizeof(MeshTree *) );

        isi = ixUpNodes.begin();

        (void) rec2RC ( &(CW->Header) , & ((*_MTree)[i].RootCell) , & nr , & nc );
        f = CW->Data[nr][nc];
        weight = (double) f ;

        for(j=0;j<nchildren;j++)
        {
            ichild = (*isi);
            ((*_MTree)[i].Child)[j] =  (*_MTree) + ichild ;
            (*_MTree)[ichild].Parent = (*_MTree) + i ;

            (void) rec2RC ( &(CW->Header) , & (((*_MTree)[i].Child)[j]->RootCell ) , & nr , & nc );
            f = CW->Data[nr][nc];

            weight = weight - (double) f ;

            isi++;
        }

        (*_MTree)[i].Weight = (float) weight ;
        meanweight += weight ;
    }

    meanweight = meanweight / nnodes ;

    // Mise a jour de la variance totale ajustable a l'amont de chaque noeud
   
    for(i=0;i<nnodes;i++)
    {
        if ( ! (*_MTree)[i].Parent )   // Si l'element n'a pas de parent
        {
            if ( MakeMeshTree_ ( (*_MTree) + i , & meanweight , & SSE ) )
            {
                SetErreur("MakeMeshTree : echec de l'appel racine ÃƒÂ  MakeMeshTree_\n" );
                return(-1);
            }

            if ( BuildPositionRange_ ( F , CW , PK , (*_MTree) + i ) )
            {
                SetErreur("MakeMeshTree : echec de l'appel racine ÃƒÂ  BuildPositionRange_\n" );
                return(-1);
            }
        }
    }

    // Reclasser les fils de chaque element par "variance totale ajustable" decroissante ?

    for(i=0;i<nnodes;i++)
    {
        MapChildren.clear();

        for(j=0;j<(*_MTree)[i].nchildren;j++)
        {
            MapChildren.insert ( make_pair ( ((*_MTree)[i].Child)[j]->SSE , ((*_MTree)[i].Child)[j] ) );
        }

        j=0;

        for(MapChildrenIterator = MapChildren.rbegin() ; MapChildrenIterator != MapChildren.rend() ; MapChildrenIterator++ )
        {
            ((*_MTree)[i].Child)[j] = (*MapChildrenIterator).second;
            j++;
        }
        
    }

    (*_meanweight) = meanweight ;

    return(0);

}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int MakeMeshTree_ ( MeshTree * MTree , double * _meanweight , double * _SSE )
{
    int i ;
    double SSEi,eps ;

    eps = (double) (MTree->Weight) - (*_meanweight) ;

/*    if ( MTree->Child[i]->cnt | MTree->nchildren > 0 )
    {
           (*_SSE) = 0 ;
    }
    else
    {*/
           (*_SSE) = eps * eps ;
//    }

    for(i=0;i<MTree->nchildren;i++)
    {
        if ( MakeMeshTree_ ( MTree->Child[i] , _meanweight , & SSEi ) )
        {
            SetErreur("MakeMeshTree_ : echec de l'appel recursif\n" );
            return(-1);
        }

//        if ( ! MTree->Child[i]->cnt )       // Si le fils Child[i] de la maille d'appel n'est pas un pt de contrainte, on propage la variance ajustable vers l'aval
//        {
            (*_SSE) += SSEi ;
//        }
    }

    MTree->SSE = (*_SSE) ;

    return(0);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int OptimMeshVariance_ ( MeshTree * MTree , double * _meanweight )
{
    int i,j,improve,pos;
    double SSEtmp,w,e,wj ;
    MeshTree * MTreeTmp , * MTreeBest ;

    if ( CopyMeshTree_ ( MTree , & MTreeBest ) )
    {
        SetErreur("OptimMeshVariance_ : echec de creation d'une copie de l'arbre\n");
        return(-1);
    }
    
    improve = 1 ;

//    while ( improve )
//    {
//        improve = 0 ;

        for(i=0;i<MTree->nchildren;i++)
        {
            //----------------------------------------------------------------------
            // Si le fils i n'est pas un point de contrainte, optimisation recursive

            if ( ! ( MTree->Child[i]->cnt ) )
            {
                if ( CopyMeshTree_ ( MTree , & MTreeTmp ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de crÃƒÂ©ation d'une copie de l'arbre\n");
                    return(-1);
                }

                if ( OptimMeshVariance_ ( MTreeTmp->Child[i] , _meanweight ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de l'appel rÃƒÂ©cursif\n");
                    return(-1);
                }

                w = (MTreeTmp->CWeight)[(MTreeTmp->CurrentPosition)];
                SSEtmp = 0 ;

                for(j=0;j<MTreeTmp->nchildren;j++)
                {
                    wj = (MTreeTmp->Child[j]->CWeight) [ MTreeTmp->Child[j]->CurrentPosition ];
                    w = w - wj ;
                    SSEtmp += (MTree->Child[j]->SSE) ;
                }

                e = w - (*_meanweight) ;
                SSEtmp += e*e ;

                MTreeTmp->Weight = (float) w;
                MTreeTmp->SSE = SSEtmp;

                if ( SSEtmp < MTreeBest->SSE )
                {
                    improve = 1 ;

                    if ( UpdateMeshTree_ ( MTreeBest , MTreeTmp ) )
                    {
                        SetErreur("OptimMeshVariance_ : echec de mise a jour de l'arbre\n");
                        return(-1);
                    }
                }

                DestroyMeshTree_ ( & MTreeTmp ) ;

            }

            //----------------------------------------------------------------------
            // Si on peut deplacer (MTree->Child)[i]->RootCell vers l'amont

            if ( test_Up ( MTree->Child[i] ) )
            {

                if ( CopyMeshTree_ ( MTree , & MTreeTmp ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de creation d'une copie de l'arbre\n");
                    return(-1);
                }

                pos = (MTreeTmp->Child[i]->CurrentPosition) ;
                pos--;
                (MTreeTmp->Child[i]->CurrentPosition) = pos;
                MTreeTmp->Child[i]->RootCell = (MTreeTmp->Child[i]->PositionRange)[pos];

                if ( OptimMeshVariance_ ( MTreeTmp->Child[i] , _meanweight ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de l'appel rÃƒÂ©cursif\n");
                    return(-1);
                }

                w = (MTreeTmp->CWeight)[(MTreeTmp->CurrentPosition)];
                SSEtmp = 0 ;

                for(j=0;j<MTreeTmp->nchildren;j++)
                {
                    wj = (MTreeTmp->Child[j]->CWeight) [ MTreeTmp->Child[j]->CurrentPosition ];
                    w = w - wj ;
                    SSEtmp += (MTree->Child[j]->SSE) ;
                }

                e = w - (*_meanweight) ;
                SSEtmp += e*e ;

                MTreeTmp->Weight = (float) w;
                MTreeTmp->SSE = SSEtmp;

                if ( SSEtmp < MTreeBest->SSE )
                {
                    improve = 1 ;

                    if ( UpdateMeshTree_ ( MTreeBest , MTreeTmp ) )
                    {
                        SetErreur("OptimMeshVariance_ : echec de mise a jour de l'arbre\n");
                        return(-1);
                    }
                }

                DestroyMeshTree_ ( & MTreeTmp ) ;
            }

            //----------------------------------------------------------------------
            // Si on peut deplacer (MTree->Child)[i]->RootCell vers l'aval

            if ( test_Down ( MTree->Child[i] ) )
            {

                if ( CopyMeshTree_ ( MTree , & MTreeTmp ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de creation d'une copie de l'arbre\n");
                    return(-1);
                }

                pos = (MTreeTmp->Child[i]->CurrentPosition) ;
                pos++;
                (MTreeTmp->Child[i]->CurrentPosition) = pos;
                MTreeTmp->Child[i]->RootCell = (MTreeTmp->Child[i]->PositionRange)[pos];

                if ( OptimMeshVariance_ ( MTreeTmp->Child[i] , _meanweight ) )
                {
                    SetErreur("OptimMeshVariance_ : echec de l'appel rÃƒÂ©cursif\n");
                    return(-1);
                }

                w = (MTreeTmp->CWeight)[(MTreeTmp->CurrentPosition)];
                SSEtmp = 0 ;

                for(j=0;j<MTreeTmp->nchildren;j++)
                {
                    wj = (MTreeTmp->Child[j]->CWeight) [ MTreeTmp->Child[j]->CurrentPosition ];
                    w = w - wj ;
                    SSEtmp += (MTree->Child[j]->SSE) ;
                }

                e = w - (*_meanweight) ;
                SSEtmp += e*e ;

                MTreeTmp->Weight = (float) w;
                MTreeTmp->SSE = SSEtmp;

                if ( SSEtmp < MTreeBest->SSE )
                {
                    improve = 1 ;

                    if ( UpdateMeshTree_ ( MTreeBest , MTreeTmp ) )
                    {
                        SetErreur("OptimMeshVariance_ : echec de mise a jour de l'arbre\n");
                        return(-1);
                    }
                }

                DestroyMeshTree_ ( & MTreeTmp ) ;
            }

        }       // Fin de la boucle sur les "MTree->nchildren"

//    }       // Fin de la boucle while
    
    //--------------------------------------------------------------------------
    // Mise a  jour finale de l'arbre d'appel avec la meilleure configuration amont

    if ( UpdateMeshTree_ ( MTree , MTreeBest ) )
    {
        SetErreur("OptimMeshVariance_ : echec de mise a jour de l'arbre\n");
        return(-1);
    }

    DestroyMeshTree_ ( & MTreeBest ) ;

    return(0);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int test_Up ( MeshTree * MTree )
{
    int upcell ;

    if ( MTree->cnt )
    {
        return(0);
    }
    
    if ( MTree->CurrentPosition == 0 )
    {
        return(0);
    }

    if (MTree->nchildren == 1)
    {
        upcell = (MTree->PositionRange)[(MTree->CurrentPosition)-1] ;

        if ( upcell == MTree->Child[0]->RootCell )
        {
            return(0);
        }
    }

    return(1);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int test_Down ( MeshTree * MTree )
{
    int downcell ;

    if ( MTree->cnt )
    {
        return(0);
    }

    if ( MTree->CurrentPosition == ((MTree->RangeLength)-1) )
    {
        return(0);
    }

    if ( MTree->Parent )
    {
        downcell = (MTree->PositionRange)[(MTree->CurrentPosition)+1] ;

        if ( downcell == MTree->Parent->RootCell )
        {
            return(0);
        }
    }
    else
    {
        return(0);
    }

    return(1);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int BuildPositionRange_ ( FlowDirGrid * F , FloatRaster * CW , FloatRaster * PK , MeshTree * MTree )
{
    int i,j,nr,nc,rec;
    int ** PathFromChild ;
    int * PathLength ;
    int * FlowLine ;
    int length;
    float f;
    set<int> TmpSet , NewTmpSet  ;
    set<int>::iterator isi ;
    map<float,int> TmpMap ;
    map<float,int>::iterator imfi ;
    MeshTree * MTreeC ;

    //--------------------------------------------------------------------------
    // Extraction des biefs reliant chaque fils a l'exutoire de la maille

    PathFromChild = (int **) malloc (  MTree->nchildren * sizeof(int *) );
    PathLength = (int *) malloc (  MTree->nchildren * sizeof(int) );

    for(i=0;i<MTree->nchildren;i++)
    {
        if ( MakeFlowLineBetweenPoints ( F , (MTree->Child)[i]->RootCell , MTree->RootCell , PathFromChild + i , PathLength + i ) )
        {
            SetErreur("BuildPositionRange_ : echec de construction du chemin %d -> %d\n",(MTree->Child)[i]->RootCell , MTree->RootCell);
            return(-1);
        }

        sort ( PathFromChild[i] , PathFromChild[i] + PathLength[i] ) ;
    }

    //--------------------------------------------------------------------------
    // Appels sur les arbres amont

    for(i=0;i<MTree->nchildren;i++)
    {
        // Appel recursif : construit la partie amont

        MTreeC = (MTree->Child)[i] ;

        BuildPositionRange_ ( F , CW , PK , MTreeC ) ;
    
        // Identification de l'ensemble des positions aval possibles du noeud (MTree->Child)[i] :
        // partie du bief (MTree->Child)[i]->RootCell ~~> MTree->RootCell  partagee avec aucun des autres fils

        TmpSet.clear();
        NewTmpSet.clear();

        TmpSet.insert ( PathFromChild[i] , PathFromChild[i] + PathLength[i] ) ;

        for(j=0;j<MTree->nchildren;j++)
        {
            if(j!=i)
            {
                set_difference ( TmpSet.begin() , TmpSet.end() , PathFromChild[j] , PathFromChild[j] + PathLength[j] , inserter(NewTmpSet,NewTmpSet.begin()) );
                TmpSet.clear();
                TmpSet.insert(NewTmpSet.begin(),NewTmpSet.end());         // Mise a jour de la difference
                NewTmpSet.clear();
            }
        }

        // Ajout des PK (set => map)

        TmpSet.erase(MTreeC->RootCell); // La position courante (racine) est deja presente dans la partie amont
        TmpMap.clear();

        for(isi=TmpSet.begin();isi!=TmpSet.end();isi++)
        {
            rec = (*isi);
            (void) rec2RC ( &(PK->Header) , & rec , & nr , & nc );
            f = PK->Data[nr][nc];
            TmpMap.insert ( make_pair ( -f , rec ) );       // classement d'amont en aval
       }

        TmpSet.clear();

        // On complete MTreeC->PositionRange avec la portion aval identifiee

        j = MTreeC->RangeLength;
        MTreeC->RangeLength += TmpMap.size();
        MTreeC->PositionRange = (int *) realloc ( MTreeC->PositionRange , (MTreeC->RangeLength) * sizeof(int) );
        MTreeC->CWeight = (double *) realloc ( MTreeC->CWeight , (MTreeC->RangeLength) * sizeof(double) );

        for(imfi=TmpMap.begin();imfi!=TmpMap.end();imfi++)
        {
            rec = (*imfi).second ;
            (MTreeC->PositionRange)[j] = rec ;

            (void) rec2RC ( &(CW->Header) , & rec , & nr , & nc );
            f = CW->Data[nr][nc];
            (MTreeC->CWeight)[j] = (double) f ;

            j++;
        }

        TmpMap.clear();
    }

    //--------------------------------------------------------------------------
    // Construction de l'intersection des "PathFromChild"

    TmpSet.clear();
    NewTmpSet.clear();

    if (MTree->nchildren > 0)
    {
        TmpSet.insert ( PathFromChild[0] , PathFromChild[0] + PathLength[0] ) ;

        for(j=1;j<MTree->nchildren;j++)
        {
            set_intersection ( TmpSet.begin() , TmpSet.end() , PathFromChild[j] , PathFromChild[j] + PathLength[j] , inserter(NewTmpSet,NewTmpSet.begin()) );
            TmpSet.clear();
            TmpSet.insert ( NewTmpSet.begin() , NewTmpSet.end() );         // Mise a jour de l'intersection
            NewTmpSet.clear();
        }
    }
    else
    {
 
        if ( MakeFlowLineThroughPoint_ ( F , CW , & (MTree->RootCell) , & FlowLine , & length ) )
        {
            SetErreur("BuildPositionRange_ : echec de construction du drain principal a l'amont de %d\n",MTree->RootCell);
            return(-1);
        }

        sort ( FlowLine , FlowLine + length ) ;
        TmpSet.insert( FlowLine , FlowLine + length );
        free(FlowLine);
    }


    if ( TmpSet.size()<=0 )
    {
        SetErreur("BuildPositionRange_ : Incoherence dans la construction des positions accessibles, la partie amont est vide\n");
        return(-1);
    }

    //--------------------------------------------------------------------------
    // Ajout des PK (set => map)

    TmpMap.clear();

    for(isi=TmpSet.begin();isi!=TmpSet.end();isi++)
    {
        rec = (*isi);
        (void) rec2RC ( &(PK->Header) , & rec , & nr , & nc );
        f = PK->Data[nr][nc];
        TmpMap.insert ( make_pair ( -f , rec ) );       // classement d'amont en aval
   }

    TmpSet.clear();

    // On initialise MTree->PositionRange avec la portion amont

    MTree->RangeLength = TmpMap.size();
    MTree->PositionRange = (int *) malloc ( (MTree->RangeLength) * sizeof(int) );
    MTree->CWeight = (double *) malloc ( (MTree->RangeLength) * sizeof(double) );
    MTree->CurrentPosition = (MTree->RangeLength) - 1;
    j=0;

    for(imfi=TmpMap.begin();imfi!=TmpMap.end();imfi++)
    {
        rec = (*imfi).second ;
        (MTree->PositionRange)[j] = rec ;
        
        (void) rec2RC ( &(CW->Header) , & rec , & nr , & nc );
         f = CW->Data[nr][nc];
        (MTree->CWeight)[j] = (double) f ;

        j++;
    }

    TmpMap.clear();

    if ( (MTree->PositionRange)[MTree->CurrentPosition] != MTree->RootCell )
    {
        SetErreur("BuildPositionRange_ : Incoherence dans la construction des positions accessibles, la position initiale n'est pas la derniere de la partie amont\n");
        return(-1);
    }

    //--------------------------------------------------------------------------
    // Desallocation des pointeurs

    for(i=0;i<MTree->nchildren;i++)
    {
        free(PathFromChild[i]);
    }
    
    free(PathFromChild);
    free(PathLength);

    return(0);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int UpdateMeshTree_ ( MeshTree * MTreeDst , MeshTree * MTreeSrc )
{
    int i;

    // Mise a  jour des elements statiques

    MTreeDst->RootCell = MTreeSrc->RootCell ;
    MTreeDst->SSE = MTreeSrc->SSE ;
    MTreeDst->Weight = MTreeSrc->Weight ;
    MTreeDst->cnt = MTreeSrc->cnt ;

    if ( MTreeDst->nchildren != MTreeSrc->nchildren )
    {
        SetErreur("UpdateMeshTree_ : les deux arbres a  refleter ont une descendance de taille differente\n");
        return(-1);
    }

    if ( MTreeDst->RangeLength != MTreeSrc->RangeLength )
    {
        SetErreur("UpdateMeshTree_ : les deux arbres a  refleter ont une plage de positions accessibles de taille differente\n");
        return(-1);
    }

    MTreeDst->CurrentPosition = MTreeSrc->CurrentPosition ;
   
    // Mise a jour recursive

    for(i=0;i<MTreeSrc->nchildren;i++)
    {
        if ( UpdateMeshTree_ ( MTreeDst->Child[i] , MTreeSrc->Child[i] ) )
        {
            SetErreur("UpdateMeshTree_ : echec de la mise a jour recursive, fils %d\n",i);
            return(-1);
        }
    }

    return(0);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int CopyMeshTree_ ( MeshTree * MTree , MeshTree ** _MTreeCopy )
{
    int i;

    if ( ! ((*_MTreeCopy) = (MeshTree *) malloc ( sizeof(MeshTree) )) )
    {
        SetErreur("CopyMeshTree_ : echec de l'allocation de l'element racine\n");
        return(-1);
    }

    // Copie des elements statiques

    (*_MTreeCopy)->RootCell = MTree->RootCell ;
    (*_MTreeCopy)->SSE = MTree->SSE ;
    (*_MTreeCopy)->Weight = MTree->Weight ;
    (*_MTreeCopy)->cnt = MTree->cnt ;
    (*_MTreeCopy)->nchildren = MTree->nchildren ;
    (*_MTreeCopy)->Parent = NULL ;

    // Copie des elements dynamiques

    (*_MTreeCopy)->CurrentPosition = MTree->CurrentPosition ;
    (*_MTreeCopy)->RangeLength = MTree->RangeLength ;
    (*_MTreeCopy)->PositionRange = (int *) malloc ( (MTree->RangeLength)*sizeof(int) );
    (*_MTreeCopy)->CWeight = (double *) malloc ( (MTree->RangeLength)*sizeof(double) );

    for (i=0;i<MTree->RangeLength;i++)
    {
        ((*_MTreeCopy)->PositionRange)[i] = (MTree->PositionRange)[i] ;
        ((*_MTreeCopy)->CWeight)[i] = (MTree->CWeight)[i] ;
    }

    // Copie recursive

    if ( !((*_MTreeCopy)->Child = (MeshTree **) malloc ( (MTree->nchildren) * sizeof ( MeshTree *) )) )
    {
        SetErreur("CopyMeshTree_ : echec de l'allocation du tableau de pointeurs vers les elements fils\n");
        return(-1);
    }

    for(i=0;i<MTree->nchildren;i++)
    {
        if ( CopyMeshTree_ ( (MTree->Child)[i] , (*_MTreeCopy)->Child + i ) )
        {
            SetErreur("CopyMeshTree_ : echec de copie recursive, fils %d\n",i);
            return(-1);
        }

        (*_MTreeCopy)->Child[i]->Parent = (*_MTreeCopy) ;
    }

    return(0);
}

//*********************************************************************************************************************************************************
//*********************************************************************************************************************************************************

int DestroyMeshTree_ ( MeshTree ** _MTree )   // destruction rÃƒÂ©cursive d'un arbre construit par CopyMeshTree_
{
    int i;

    free ( (*_MTree)->PositionRange );
    free ( (*_MTree)->CWeight );

    for(i=0;i<(*_MTree)->nchildren;i++)
    {
        DestroyMeshTree_ ( (*_MTree)->Child + i );
    }

    free((*_MTree)->Child);
    free(*_MTree);

    return(0);

}

