/* ======================================================================== *\
!
! *
! * This file is part of MARS, the MAGIC Analysis and Reconstruction
! * Software. It is distributed to you in the hope that it can be a useful
! * and timesaving tool in analysing Data of imaging Cerenkov telescopes.
! * It is distributed WITHOUT ANY WARRANTY.
! *
! * Permission to use, copy, modify and distribute this software and its
! * documentation for any purpose is hereby granted without fee,
! * provided that the above copyright notice appear in all copies and
! * that both that copyright notice and this permission notice appear
! * in supporting documentation. It is provided "as is" without express
! * or implied warranty.
! *
!
!
!   Author(s): Thomas Hengstebeck 3/2003 <mailto:hengsteb@alwa02.physik.uni-siegen.de>
!
!   Copyright: MAGIC Software Development, 2000-2003
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//                                                                         //
// MRanTree                                                                //
//                                                                         //
// ParameterContainer for Tree structure                                   //
//                                                                         //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////
#include "MRanTree.h"

#include <iostream>

#include <TVector.h>
#include <TMatrix.h>
#include <TRandom.h>

#include "MDataArray.h"

#include "MLog.h"
#include "MLogManip.h"

ClassImp(MRanTree);

using namespace std;

// --------------------------------------------------------------------------
//
// Default constructor.
//
MRanTree::MRanTree(const char *name, const char *title):fNdSize(0), fNumTry(3), fData(NULL)
{

    fName  = name  ? name  : "MRanTree";
    fTitle = title ? title : "Storage container for structure of a single tree";
}

void MRanTree::SetNdSize(Int_t n)
{
    // threshold nodesize of terminal nodes, i.e. the training data is splitted
    // until there is only pure date in the subsets(=terminal nodes) or the
    // subset size is LE n

    fNdSize=TMath::Max(1,n);//at least 1 event per node
}

void MRanTree::SetNumTry(Int_t n)
{
    // number of trials in random split selection:
    // choose at least 1 variable to split in

    fNumTry=TMath::Max(1,n);
}

void MRanTree::GrowTree(const TMatrix &mhad, const TMatrix &mgam,
                        const TArrayI &hadtrue, TArrayI &datasort,
                        const TArrayI &datarang, TArrayF &tclasspop, TArrayI &jinbag,
                        const TArrayF &winbag)
{
    // arrays have to be initialized with generous size, so number of total nodes (nrnodes)
    // is estimated for worst case
    const Int_t numdim =mhad.GetNcols();
    const Int_t numdata=winbag.GetSize();
    const Int_t nrnodes=2*numdata+1;

    // number of events in bootstrap sample
    Int_t ninbag=0;
    for (Int_t n=0;n<numdata;n++)
        if(jinbag[n]==1) ninbag++;

    TArrayI bestsplit(nrnodes);
    TArrayI bestsplitnext(nrnodes);

    fBestVar.Set(nrnodes);
    fTreeMap1.Set(nrnodes);
    fTreeMap2.Set(nrnodes);
    fBestSplit.Set(nrnodes);

    fTreeMap1.Reset();
    fTreeMap2.Reset();
    fBestSplit.Reset();

    fGiniDec.Set(numdim);
    fGiniDec.Reset();

    // tree growing
    BuildTree(datasort,datarang,hadtrue,bestsplit,
              bestsplitnext,tclasspop,winbag,ninbag);

    // post processing, determine cut (or split) values fBestSplit
    Int_t nhad=mhad.GetNrows();

    for(Int_t k=0; k<nrnodes; k++)
    {
        if (GetNodeStatus(k)==-1)
            continue;

        const Int_t &bsp =bestsplit[k];
        const Int_t &bspn=bestsplitnext[k];
        const Int_t &msp =fBestVar[k];

        fBestSplit[k]  = bsp<nhad  ? mhad(bsp, msp):mgam(bsp-nhad, msp);
        fBestSplit[k] += bspn<nhad ? mhad(bspn,msp):mgam(bspn-nhad,msp);
        fBestSplit[k] /= 2;
    }

    // resizing arrays to save memory
    fBestVar.Set(fNumNodes);
    fTreeMap1.Set(fNumNodes);
    fTreeMap2.Set(fNumNodes);
    fBestSplit.Set(fNumNodes);
}

Int_t MRanTree::FindBestSplit(const TArrayI &datasort,const TArrayI &datarang,
                              const TArrayI &hadtrue,Int_t ndstart,Int_t ndend,TArrayF &tclasspop,
                              Int_t &msplit,Float_t &decsplit,Int_t &nbest,
                              const TArrayF &winbag)
{
    const Int_t nrnodes = fBestSplit.GetSize();
    const Int_t numdata = (nrnodes-1)/2;
    const Int_t mdim = fGiniDec.GetSize();

    // weighted class populations after split
    TArrayF wc(2); 
    TArrayF wr(2); // right node

    // For the best split, msplit is the index of the variable (e.g Hillas par., zenith angle ,...)
    // split on. decsplit is the decreae in impurity measured by Gini-index.
    // nsplit is the case number of value of msplit split on,
    // and nsplitnext is the case number of the next larger value of msplit.

    Int_t nbestvar=0;

    // compute initial values of numerator and denominator of Gini-index,
    // Gini index= pno/dno
    Double_t pno=0;
    Double_t pdo=0;
    for (Int_t j=0;j<2;j++)
    {
        pno+=tclasspop[j]*tclasspop[j];
        pdo+=tclasspop[j];
    }

    const Double_t crit0=pno/pdo;
    Int_t jstat=0;

    // start main loop through variables to find best split,
    // (Gini-index as criterium crit)

    Double_t critmax=-1.0e20;  // FIXME: Replace by a constant from limits.h

    // random split selection, number of trials = fNumTry
    for(Int_t mt=0;mt<fNumTry;mt++)
    {
        const Int_t mvar=Int_t(gRandom->Rndm()*mdim);
        const Int_t mn  = mvar*numdata;

        // Gini index = rrn/rrd+rln/rld
        Double_t rrn=pno;
        Double_t rrd=pdo;
        Double_t rln=0;
        Double_t rld=0;

        TArrayF wl(2); // left node
        wr = tclasspop;

        Double_t critvar=-1.0e20;

        for(Int_t nsp=ndstart;nsp<=ndend-1;nsp++)
        {
            const Int_t &nc=datasort[mn+nsp];
            const Int_t &k=hadtrue[nc];

            const Float_t &u=winbag[nc];

            rln+=u*(2*wl[k]+u);
            rrn+=u*(-2*wr[k]+u);
            rld+=u;
            rrd-=u;

            wl[k]+=u;
            wr[k]-=u;

            if (datarang[mn+nc]>=datarang[mn+datasort[mn+nsp+1]])
                continue;
            if (TMath::Min(rrd,rld)<=1.0e-5)
                continue;

            const Double_t crit=(rln/rld)+(rrn/rrd);
            if (crit<=critvar)
                continue;

            nbestvar=nsp;
            critvar=crit;
        }

        if (critvar<=critmax)
            continue;

        msplit=mvar;
        nbest=nbestvar;
        critmax=critvar;
    }

    decsplit=critmax-crit0;

    return critmax<-1.0e10 ? 1 : jstat;
}

void MRanTree::MoveData(TArrayI &datasort,Int_t ndstart,
                        Int_t ndend,TArrayI &idmove,TArrayI &ncase,Int_t msplit,
                        Int_t nbest,Int_t &ndendl)
{
    // This is the heart of the BuildTree construction. Based on the best split
    // the data in the part of datasort corresponding to the current node is moved to the
    // left if it belongs to the left child and right if it belongs to the right child-node.
    const Int_t numdata = ncase.GetSize();
    const Int_t mdim    = fGiniDec.GetSize();

    TArrayI tdatasort(numdata);

    // compute idmove = indicator of case nos. going left

    for (Int_t nsp=ndstart;nsp<=ndend;nsp++)
    {
        const Int_t &nc=datasort[msplit*numdata+nsp];
        idmove[nc]= nsp<=nbest?1:0;
    }
    ndendl=nbest;

    // shift case. nos. right and left for numerical variables.

    for(Int_t msh=0;msh<mdim;msh++)
    {
        Int_t k=ndstart-1;
        for (Int_t n=ndstart;n<=ndend;n++)
        {
            const Int_t &ih=datasort[msh*numdata+n];
            if (idmove[ih]==1)
                tdatasort[++k]=datasort[msh*numdata+n];
        }

        for (Int_t n=ndstart;n<=ndend;n++)
        {
            const Int_t &ih=datasort[msh*numdata+n];
            if (idmove[ih]==0)
                tdatasort[++k]=datasort[msh*numdata+n];
        }

        for(Int_t m=ndstart;m<=ndend;m++)
            datasort[msh*numdata+m]=tdatasort[m];
    }

    // compute case nos. for right and left nodes.

    for(Int_t n=ndstart;n<=ndend;n++)
        ncase[n]=datasort[msplit*numdata+n];
}

void MRanTree::BuildTree(TArrayI &datasort,const TArrayI &datarang,
                         const TArrayI &hadtrue, TArrayI &bestsplit,
                         TArrayI &bestsplitnext, TArrayF &tclasspop,
                         const TArrayF &winbag, Int_t ninbag)
{
    // Buildtree consists of repeated calls to two void functions, FindBestSplit and MoveData.
    // Findbestsplit does just that--it finds the best split of the current node.
    // MoveData moves the data in the split node right and left so that the data
    // corresponding to each child node is contiguous.
    //
    // buildtree bookkeeping:
    // ncur is the total number of nodes to date.  nodestatus(k)=1 if the kth node has been split.
    // nodestatus(k)=2 if the node exists but has not yet been split, and =-1 if the node is
    // terminal.  A node is terminal if its size is below a threshold value, or if it is all
    // one class, or if all the data-values are equal.  If the current node k is split, then its
    // children are numbered ncur+1 (left), and ncur+2(right), ncur increases to ncur+2 and
    // the next node to be split is numbered k+1.  When no more nodes can be split, buildtree
    // returns.
    const Int_t mdim    = fGiniDec.GetSize();
    const Int_t nrnodes = fBestSplit.GetSize();
    const Int_t numdata = (nrnodes-1)/2;

    TArrayI nodepop(nrnodes);
    TArrayI nodestart(nrnodes);
    TArrayI parent(nrnodes);

    TArrayI ncase(numdata);
    TArrayI idmove(numdata);
    TArrayI iv(mdim);

    TArrayF classpop(nrnodes*2);
    TArrayI nodestatus(nrnodes);

    for (Int_t j=0;j<2;j++)
        classpop[j*nrnodes+0]=tclasspop[j];

    Int_t ncur=0;
    nodepop[0]=ninbag;
    nodestatus[0]=2;

    // start main loop
    for (Int_t kbuild=0; kbuild<nrnodes; kbuild++)
    {
          if (kbuild>ncur) break;
          if (nodestatus[kbuild]!=2) continue;

          // initialize for next call to FindBestSplit

          const Int_t ndstart=nodestart[kbuild];
          const Int_t ndend=ndstart+nodepop[kbuild]-1;
          for (Int_t j=0;j<2;j++)
              tclasspop[j]=classpop[j*nrnodes+kbuild];

          Int_t msplit, nbest;
          Float_t decsplit=0;
          const Int_t jstat=FindBestSplit(datasort,datarang,hadtrue,
                                          ndstart,ndend,tclasspop,msplit,
                                          decsplit,nbest,winbag);

          if (jstat==1)
          {
              nodestatus[kbuild]=-1;
              continue;
          }

          fBestVar[kbuild]=msplit;
          fGiniDec[msplit]+=decsplit;

          bestsplit[kbuild]=datasort[msplit*numdata+nbest];
          bestsplitnext[kbuild]=datasort[msplit*numdata+nbest+1];

          Int_t ndendl;
          MoveData(datasort,ndstart,ndend,idmove,ncase,
                   msplit,nbest,ndendl);

          // leftnode no.= ncur+1, rightnode no. = ncur+2.

          nodepop[ncur+1]=ndendl-ndstart+1;
          nodepop[ncur+2]=ndend-ndendl;
          nodestart[ncur+1]=ndstart;
          nodestart[ncur+2]=ndendl+1;

          // find class populations in both nodes

          for (Int_t n=ndstart;n<=ndendl;n++)
          {
              const Int_t &nc=ncase[n];
              const Int_t &j=hadtrue[nc];
              classpop[j*nrnodes+ncur+1]+=winbag[nc];
          }

          for (Int_t n=ndendl+1;n<=ndend;n++)
          {
              const Int_t &nc=ncase[n];
              const Int_t &j=hadtrue[nc];
              classpop[j*nrnodes+ncur+2]+=winbag[nc];
          }

          // check on nodestatus

          nodestatus[ncur+1]=2;
          nodestatus[ncur+2]=2;
          if (nodepop[ncur+1]<=fNdSize) nodestatus[ncur+1]=-1;
          if (nodepop[ncur+2]<=fNdSize) nodestatus[ncur+2]=-1;

          Double_t popt1=0;
          Double_t popt2=0;
          for (Int_t j=0;j<2;j++)
          {
              popt1+=classpop[j*nrnodes+ncur+1];
              popt2+=classpop[j*nrnodes+ncur+2];
          }

          for (Int_t j=0;j<2;j++)
          {
              if (classpop[j*nrnodes+ncur+1]==popt1) nodestatus[ncur+1]=-1;
              if (classpop[j*nrnodes+ncur+2]==popt2) nodestatus[ncur+2]=-1;
          }

          fTreeMap1[kbuild]=ncur+1;
          fTreeMap2[kbuild]=ncur+2;
          parent[ncur+1]=kbuild;
          parent[ncur+2]=kbuild;
          nodestatus[kbuild]=1;
          ncur+=2;
          if (ncur>=nrnodes) break;
    }

    // determine number of nodes
    fNumNodes=nrnodes;
    for (Int_t k=nrnodes-1;k>=0;k--)
    {
        if (nodestatus[k]==0) fNumNodes-=1;
        if (nodestatus[k]==2) nodestatus[k]=-1;
    }

    fNumEndNodes=0;
    for (Int_t kn=0;kn<fNumNodes;kn++)
        if(nodestatus[kn]==-1)
        {
            fNumEndNodes++;
            Double_t pp=0;
            for (Int_t j=0;j<2;j++)
            {
                if(classpop[j*nrnodes+kn]>pp)
                {
                    // class + status of node kn coded into fBestVar[kn]
                    fBestVar[kn]=j-2;
                    pp=classpop[j*nrnodes+kn];
                }
            }
            fBestSplit[kn] =classpop[1*nrnodes+kn];
            fBestSplit[kn]/=(classpop[0*nrnodes+kn]+classpop[1*nrnodes+kn]);
        }
}

void MRanTree::SetRules(MDataArray *rules)
{
    fData=rules;
}

Double_t MRanTree::TreeHad(const TVector &event)
{
    Int_t kt=0;
    // to optimize on storage space node status and node class
    // are coded into fBestVar:
    // status of node kt = TMath::Sign(1,fBestVar[kt])
    // class  of node kt = fBestVar[kt]+2 (class defined by larger
    //  node population, actually not used)
    // hadronness assigned to node kt = fBestSplit[kt]

    for (Int_t k=0;k<fNumNodes;k++)
    {
        if (fBestVar[kt]<0)
            break;

        const Int_t m=fBestVar[kt];
        kt = event(m)<=fBestSplit[kt] ? fTreeMap1[kt] : fTreeMap2[kt];
    }

    return fBestSplit[kt];
}

Double_t MRanTree::TreeHad(const TMatrixRow &event)
{
    Int_t kt=0;
    // to optimize on storage space node status and node class
    // are coded into fBestVar:
    // status of node kt = TMath::Sign(1,fBestVar[kt])
    // class  of node kt = fBestVar[kt]+2 (class defined by larger
    //  node population, actually not used)
    // hadronness assigned to node kt = fBestSplit[kt]

    for (Int_t k=0;k<fNumNodes;k++)
    {
        if (fBestVar[kt]<0)
            break;

        const Int_t m=fBestVar[kt];
        kt = event(m)<=fBestSplit[kt] ? fTreeMap1[kt] : fTreeMap2[kt];
    }

    return fBestSplit[kt];
}

Double_t MRanTree::TreeHad(const TMatrix &m, Int_t ievt)
{
    return TreeHad(TMatrixRow(m, ievt));
}

Double_t MRanTree::TreeHad()
{
    TVector event;
    *fData >> event;

    return TreeHad(event);
}

Bool_t MRanTree::AsciiWrite(ostream &out) const
{
    TString str;
    Int_t k;

    out.width(5);out<<fNumNodes<<endl;

    for (k=0;k<fNumNodes;k++)
    {
        str=Form("%f",GetBestSplit(k));

        out.width(5);  out << k;
        out.width(5);  out << GetNodeStatus(k);
        out.width(5);  out << GetTreeMap1(k);
        out.width(5);  out << GetTreeMap2(k);
        out.width(5);  out << GetBestVar(k);
        out.width(15); out << str<<endl;
        out.width(5);  out << GetNodeClass(k);
    }
    out<<endl;

    return k==fNumNodes;
}