/* ======================================================================== *\
!
! *
! * 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 2/2005 <mailto:hengsteb@physik.hu-berlin.de>
!
!   Copyright: MAGIC Software Development, 2000-2005
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//
//  MRFEnergyEst
//
//
////////////////////////////////////////////////////////////////////////////
#include "MRFEnergyEst.h"

#include <TFile.h>
#include <TList.h>

#include <TH1.h>
#include <TH2.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <TMath.h>
#include <TVector.h>

#include "MHMatrix.h"

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

#include "MParList.h"
#include "MTaskList.h"
#include "MEvtLoop.h"

#include "MRanTree.h"
#include "MRanForest.h"
#include "MRanForestGrow.h"

#include "MData.h"
#include "MParameters.h"

ClassImp(MRFEnergyEst);

using namespace std;

static const TString gsDefName  = "MRFEnergyEst";
static const TString gsDefTitle = "RF for energy estimation";

// --------------------------------------------------------------------------
//
//  Default constructor. Set the name and title of this object
//
MRFEnergyEst::MRFEnergyEst(const char *name, const char *title)
    : fDebug(kFALSE), fData(0), fEnergyEst(0),
    fNumTrees(-1), fNumTry(-1), fNdSize(-1), 
    fTestMatrix(0), fEstimationMode(kMean)
{
    fName  = name  ? name  : gsDefName.Data();
    fTitle = title ? title : gsDefTitle.Data();
}

MRFEnergyEst::~MRFEnergyEst()
{
    fEForests.Delete();
}

// --------------------------------------------------------------------------
//
// Train the RF with the goven MHMatrix. The last column must contain the
// True energy.
//
Int_t MRFEnergyEst::Train(const MHMatrix &matrixtrain, const TArrayD &grid)
{
    gLog.Separator("MRFEnergyEst - Train");

    if (!matrixtrain.GetColumns())
    {
        *fLog << err << "ERROR - MHMatrix does not contain rules... abort." << endl;
        return kFALSE;
    }

    const Int_t ncols = matrixtrain.GetM().GetNcols();
    const Int_t nrows = matrixtrain.GetM().GetNrows();
    if (ncols<=0 || nrows <=0)
    {
        *fLog << err << "ERROR - No. of columns or no. of rows of matrixtrain equal 0 ... abort." << endl;
        return kFALSE;
    }

    const Int_t nbins = grid.GetSize()-1;
    if (nbins<=0)
    {
        *fLog << err << "ERROR - Energy grid not vaild... abort." << endl;
        return kFALSE;
    }

    // rules (= combination of image par) to be used for energy estimation
    TFile fileRF(fFileName, "recreate");
    if (!fileRF.IsOpen())
    {
        *fLog << err << "ERROR - File to store RFs could not be opened... abort." << endl;
        return kFALSE;
    }

    const Int_t nobs = 3; // Number of obsolete columns

    MDataArray usedrules;
    for(Int_t i=0; i<ncols-nobs; i++) // -3 is important!!!
        usedrules.AddEntry((*matrixtrain.GetColumns())[i].GetRule());

    // training of RF for each energy bin
    for (Int_t ie=0; ie<nbins; ie++)
    {
        TMatrix mat1(nrows, ncols);
        TMatrix mat0(nrows, ncols);

        // prepare matrix for current energy bin
        Int_t irow1=0;
        Int_t irow0=0;

        const TMatrix &m = matrixtrain.GetM();
        for (Int_t j=0; j<nrows; j++)
        {
            const Double_t energy = m(j,ncols-1);

            if (energy>grid[ie] && energy<=grid[ie+1])
                TMatrixFRow(mat1, irow1++) = TMatrixFRow_const(m,j);
            else
                TMatrixFRow(mat0, irow0++) = TMatrixFRow_const(m,j);
        }

        const Bool_t invalid = irow1==0 || irow0==0;

        if (invalid)
            *fLog << warn << "WARNING - Skipping";
        else
            *fLog << inf << "Training RF for";

        *fLog << " energy bin " << ie << " (" << grid[ie] << ", " << grid[ie+1] << ") " << irow0 << " " << irow1 << endl;

        if (invalid)
            continue;

        if (fDebug)
            gLog.SetNullOutput(kTRUE);

        // last column must be removed (true energy col.)
        mat1.ResizeTo(irow1, ncols-nobs);
        mat0.ResizeTo(irow0, ncols-nobs);

        // create MHMatrix as input for RF
        MHMatrix matrix1(mat1, "MatrixHadrons");
        MHMatrix matrix0(mat0, "MatrixGammas");

        // training of RF
        MTaskList tlist;
        MParList plist;
        plist.AddToList(&tlist);
        plist.AddToList(&matrix0);
        plist.AddToList(&matrix1);

        MRanForestGrow rfgrow;
        rfgrow.SetNumTrees(fNumTrees); // number of trees
        rfgrow.SetNumTry(fNumTry);     // number of trials in random split selection
        rfgrow.SetNdSize(fNdSize);     // limit for nodesize

        tlist.AddToList(&rfgrow);
    
        MEvtLoop evtloop;
        evtloop.SetDisplay(fDisplay);
        evtloop.SetParList(&plist);

        if (!evtloop.Eventloop())
            return kFALSE;

        if (fDebug)
            gLog.SetNullOutput(kFALSE);

        // Calculate bin center
        const Double_t E = (TMath::Log10(grid[ie])+TMath::Log10(grid[ie+1]))/2;

        // save whole forest
        MRanForest *forest=(MRanForest*)plist.FindObject("MRanForest");
        forest->SetUserVal(E);
        forest->Write(Form("%.10f", E));
    }

    // save rules
    usedrules.Write("rules");

    return kTRUE;
}

Int_t MRFEnergyEst::ReadForests(MParList &plist)
{
    TFile fileRF(fFileName, "read");
    if (!fileRF.IsOpen())
    {
        *fLog << err << dbginf << "File containing RFs could not be opened... aborting." << endl;
        return kFALSE;
    }

    fEForests.Delete();

    TIter Next(fileRF.GetListOfKeys());
    TObject *o=0;
    while ((o=Next()))
    {
        MRanForest *forest;
        fileRF.GetObject(o->GetName(), forest);
        if (!forest)
            continue;

        forest->SetUserVal(atof(o->GetName()));

        fEForests.Add(forest);
    }

    if (fData->Read("rules")<=0)
    {
        *fLog << err << "ERROR - Reading 'rules' from file " << fFileName << endl;
        return kFALSE;
    }

    return kTRUE;
}

Int_t MRFEnergyEst::PreProcess(MParList *plist)
{
    fEnergyEst = (MParameterD*)plist->FindCreateObj("MParameterD", "MEnergyEst");
    if (!fEnergyEst)
        return kFALSE;

    fData = (MDataArray*)plist->FindCreateObj("MDataArray");
    if (!fData)
        return kFALSE;

    if (!ReadForests(*plist))
    {
        *fLog << err << "Reading RFs failed... aborting." << endl;
        return kFALSE;
    }

    *fLog << inf << "RF read from " << fFileName << endl;

    if (fTestMatrix)
        return kTRUE;

    fData->Print();

    if (!fData->PreProcess(plist))
    {
        *fLog << err << "PreProcessing of the MDataArray failed... aborting." << endl;
        return kFALSE;
    }

    return kTRUE;
}

// --------------------------------------------------------------------------
//
//
#include <TGraph.h>
#include <TF1.h>
Int_t MRFEnergyEst::Process()
{
    static TF1 f1("f1", "gaus");

    TVector event;
    if (fTestMatrix)
        *fTestMatrix >> event;
    else
        *fData >> event;

    Double_t sume = 0;
    Double_t sumh = 0;
    Double_t maxh = 0;
    Double_t maxe = 0;

    Double_t max  = -1e10;
    Double_t min  =  1e10;

    //TH1C h("", "", fEForests.GetSize(), 0, 1);

    TIter Next(&fEForests);
    MRanForest *rf = 0;

    TGraph g;
    while ((rf=(MRanForest*)Next()))
    {
        const Double_t h = rf->CalcHadroness(event);
        const Double_t e = rf->GetUserVal();
        g.SetPoint(g.GetN(), e, h);
        sume += e*h;
        sumh += h;
        if (h>maxh)
        {
            maxh = h;
            maxe = e;
        }
        if (e>max)
            max = e;
        if (e<min)
            min = e;
    }

    switch (fEstimationMode)
    {
    case kMean:
        fEnergyEst->SetVal(pow(10, sume/sumh));
        break;
    case kMaximum:
        fEnergyEst->SetVal(pow(10, maxe));
        break;
    case kFit:
        f1.SetParameter(0, maxh);
        f1.SetParameter(1, maxe);
        f1.SetParameter(2, 0.125);
        g.Fit(&f1, "Q0N");
        fEnergyEst->SetVal(pow(10, f1.GetParameter(1)));
        break;

    }
    fEnergyEst->SetReadyToSave();

    return kTRUE;
}

// --------------------------------------------------------------------------
//
//
Int_t MRFEnergyEst::ReadEnv(const TEnv &env, TString prefix, Bool_t print)
{
    Bool_t rc = kFALSE;
    if (IsEnvDefined(env, prefix, "FileName", print))
    {
        rc = kTRUE;
        SetFileName(GetEnvValue(env, prefix, "FileName", fFileName));
    }
    if (IsEnvDefined(env, prefix, "Debug", print))
    {
        rc = kTRUE;
        SetDebug(GetEnvValue(env, prefix, "Debug", fDebug));
    }
    if (IsEnvDefined(env, prefix, "EstimationMode", print))
    {
        TString txt = GetEnvValue(env, prefix, "EstimationMode", "");
        txt = txt.Strip(TString::kBoth);
        txt.ToLower();
        if (txt==(TString)"mean")
            fEstimationMode = kMean;
        if (txt==(TString)"maximum")
            fEstimationMode = kMaximum;
        if (txt==(TString)"fit")
            fEstimationMode = kFit;
        rc = kTRUE;
    }
    return rc;
}