/* ======================================================================== *\
!
! *
! * 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 Bretz 11/2005 <mailto:tbretz@astro.uni-wuerzburg.de>
!
!   Copyright: MAGIC Software Development, 2005
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//
// MJTrainDisp
//
//
// Example:
// --------
//
//    // SequencesOn are used for training, SequencesOff for Testing
//    MDataSet set("mctesttrain.txt");
//    set.SetNumAnalysis(1);  // Must have a number
//    MJTrainDisp opt;
//    //opt.SetDebug();
//    opt.AddParameter("MHillas.fLength");
//    opt.AddParameter("MHillas.fWidth");
//    MStatusDisplay *d = new MStatusDisplay;
//    opt.SetDisplay(d);
//    opt.AddPreCut("MHillasSrc.fDCA*MGeomCam.fConvMm2Deg<0.3");
//    opt.Train("rf-disp.root", set, 30000); // Number of train events
//
//
// Random Numbers:
// ---------------
//   Use:
//         if(gRandom)
//             delete gRandom;
//         gRandom = new TRandom3();
//   in advance to change the random number generator.
//
////////////////////////////////////////////////////////////////////////////
#include "MJTrainDisp.h"

#include "MHMatrix.h"

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

// tools
#include "MDataSet.h"
#include "MTFillMatrix.h"
#include "MChisqEval.h"

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

// tasks
#include "MReadMarsFile.h"
#include "MContinue.h"
#include "MFillH.h"
#include "MRanForestCalc.h"
#include "MParameterCalc.h"

// container
#include "MParameters.h"

// histograms
#include "MBinning.h"
#include "MH3.h"
#include "MHThetaSq.h"

// filter
#include "MFilterList.h"

ClassImp(MJTrainDisp);

using namespace std;

Bool_t MJTrainDisp::Train(const char *out, const MDataSet &set, Int_t num)
{
    if (!set.IsValid())
    {
        *fLog << err << "ERROR - DataSet invalid!" << endl;
        return kFALSE;
    }

    *fLog << inf;
    fLog->Separator(GetDescriptor());

    // --------------------- Setup files --------------------
    MReadMarsFile readtrn("Events");
    MReadMarsFile readtst("Events");
    readtrn.DisableAutoScheme();
    readtst.DisableAutoScheme();

    if (!set.AddFilesOn(readtrn))
        return kFALSE;
    if (!set.AddFilesOff(readtst))
        return kFALSE;

    // ----------------------- Setup Matrix ------------------
    MHMatrix train("Train");
    train.AddColumns(fRules);
    if (fEnableWeights)
        train.AddColumn("MWeight.fVal");
    train.AddColumn("MHillasSrc.fDist*MGeomCam.fConvMm2Deg");
    //train.AddColumn("TMath::Hypot(MHillasSrc.fDCA, MHillasSrc.fDist)*MGeomCam.fConvMm2Deg");

    // ----------------------- Fill Matrix RF ----------------------
    MTFillMatrix fill;
    fill.SetDisplay(fDisplay);
    fill.SetLogStream(fLog);
    fill.SetDestMatrix1(&train, num);
    fill.SetReader(&readtrn);
    fill.AddPreCuts(fPreCuts);
    fill.AddPreCuts(fTrainCuts);
    fill.AddPreTasks(fPreTasks);
    fill.AddPostTasks(fPostTasks);
    if (!fill.Process())
        return kFALSE;

    // ------------------------ Train RF --------------------------
    MRanForestCalc rf;
    rf.SetNumTrees(fNumTrees);
    rf.SetNdSize(fNdSize);
    rf.SetNumTry(fNumTry);
    rf.SetNumObsoleteVariables(1);
    rf.SetLastDataColumnHasWeights(fEnableWeights);
    rf.SetDisplay(fDisplay);
    rf.SetLogStream(fLog);
    rf.SetFileName(out);
    rf.SetDebug(fDebug);
    rf.SetNameOutput("Disp");

    /*
    MBinning b(32, 10, 100000, "BinningEnergyEst", "log");

    if (!rf.TrainMultiRF(train, b.GetEdgesD()))    // classification with one tree per bin
        return;

    if (!rf.TrainSingleRF(train, b.GetEdgesD()))   // classification into different bins
        return;
    */
    if (!rf.TrainRegression(train))                  // regression (best choice)
        return kFALSE;

    // --------------------- Display result ----------------------

    gLog.Separator("Test");

    MParList  plist;
    MTaskList tlist;
    plist.AddToList(this);
    plist.AddToList(&tlist);
    //plist.AddToList(&b);

    MParameterD par("ThetaSquaredCut");
    par.SetVal(0.2);
    plist.AddToList(&par);

    MAlphaFitter fit;
    fit.SetPolynomOrder(0);
    fit.SetSignalFitMax(0.8);
    fit.EnableBackgroundFit(kFALSE);
    fit.SetSignalFunction(MAlphaFitter::kThetaSq);
    fit.SetMinimizationStrategy(MAlphaFitter::kGaussSigma);
    plist.AddToList(&fit);

    MFilterList list;
    if (!list.AddToList(fPreCuts))
        *fLog << err << "ERROR - Calling MFilterList::AddToList for fPreCuts failed!" << endl;
    if (!list.AddToList(fTestCuts))
        *fLog << err << "ERROR - Calling MFilterList::AddToList for fTestCuts failed!" << endl;

    MContinue cont(&list);
    cont.SetInverted();

    MHThetaSq hist;
    hist.SkipHistTime();
    hist.SkipHistTheta();
    hist.SkipHistEnergy();

    MFillH fillh(&hist, "", "FillThetaSq");

    // 0 =  disp^2 - 2*disp*dist*cos(alpha) + dist^2

    // cos^2 -1 = - sin^2

    // disp = +dist* (cos(alpha) +/- sqrt(cos^2(alpha) - 1) )

    const char *rule = "(MHillasSrc.fDist*MGeomCam.fConvMm2Deg)^2 + (Disp.fVal)^2 - (2*MHillasSrc.fDist*MGeomCam.fConvMm2Deg*Disp.fVal*cos(MHillasSrc.fAlpha*kDeg2Rad))";

    MParameterCalc calcthetasq(rule, "MThetaSqCalc");
    calcthetasq.SetNameParameter("ThetaSquared");

    MChisqEval eval;
    eval.SetY1("sqrt(ThetaSquared.fVal)");

    MH3 hdisp1("MHillas.fSize",  "sqrt(ThetaSquared.fVal)");
    MH3 hdisp2("MMcEvt.fEnergy", "sqrt(ThetaSquared.fVal)");
    hdisp1.SetTitle("\\vartheta distribution vs. Size:Size [phe]:\\vartheta [\\circ]");
    hdisp2.SetTitle("\\vartheta distribution vs. Energy:Enerhy [GeV]:\\vartheta [\\circ]");

    MBinning binsx(50, 10, 100000, "BinningMH3X", "log");
    MBinning binsy(50, 0,  1,      "BinningMH3Y", "lin");

    plist.AddToList(&binsx);
    plist.AddToList(&binsy);

    MFillH fillh2a(&hdisp1, "", "FillSize");
    MFillH fillh2b(&hdisp2, "", "FillEnergy");
    fillh2a.SetDrawOption("blue profx");
    fillh2b.SetDrawOption("blue profx");
    fillh2a.SetNameTab("Size");
    fillh2b.SetNameTab("Energy");

    tlist.AddToList(&readtst);
    tlist.AddToList(&cont);
    tlist.AddToList(&rf);
    tlist.AddToList(&calcthetasq);
    tlist.AddToList(&fillh);
    tlist.AddToList(&fillh2a);
    tlist.AddToList(&fillh2b);
    tlist.AddToList(&eval);

    MEvtLoop loop;
    loop.SetLogStream(fLog);
    loop.SetDisplay(fDisplay);
    loop.SetParList(&plist);

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

    // Print the result
    *fLog << inf << "Rule: " << rule << endl;
    hist.GetAlphaFitter().Print("result");

    if (!WriteDisplay(out))
        return kFALSE;

    return kTRUE;
}