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

/////////////////////////////////////////////////////////////////////////////
//
//  MJCut
//
// FIXME: Preparation for wobble mode missing
//
/////////////////////////////////////////////////////////////////////////////
#include "MJCut.h"

#include <TEnv.h>
#include <TFile.h>

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

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

#include "MStatusDisplay.h"

#include "MReadReports.h"
#include "MPrint.h"
#include "MContinue.h"
#include "MEnergyEstimate.h"
#include "MTaskEnv.h"
#include "MSrcPosCalc.h"
#include "MHillasCalc.h"
#include "MFillH.h"
#include "MWriteRootFile.h"

#include "../mhflux/MAlphaFitter.h"
#include "MH3.h"
#include "MBinning.h"
#include "MDataSet.h"
#include "MParameters.h"
#include "MPointingPos.h"
#include "MObservatory.h"

ClassImp(MJCut);

using namespace std;

// --------------------------------------------------------------------------
//
// Default constructor.  Set defaults for fStoreSummary, fStoreresult,
// fWriteOnly, fIsWobble and fFullDisplay to kFALSE and initialize
// fEstimateEnergy and fCalcHadronness with NULL.
//
MJCut::MJCut(const char *name, const char *title)
    : fStoreSummary(kFALSE), fStoreResult(kFALSE), fWriteOnly(kFALSE),
    fIsWobble(kFALSE), fFullDisplay(kFALSE), /*fSubstraction(kFALSE),*/
    fEstimateEnergy(0), fCalcHadronness(0)
{
    fName  = name  ? name  : "MJCut";
    fTitle = title ? title : "Standard program to perform g/h-separation cuts";
}

// --------------------------------------------------------------------------
//
// Destructor. Delete fEstimateEnergy and fCalcHadronness if != NULL
//
MJCut::~MJCut()
{
    if (fEstimateEnergy)
        delete fEstimateEnergy;
    if (fCalcHadronness)
        delete fCalcHadronness;
}

// --------------------------------------------------------------------------
//
// Set the name of the summary file (events after cut0)
// If you give a name the storage of this file is enabled implicitly.
// If you give no filename the storage is neither enabled nor disabled,
// but the storage file name is reset.
// If no filename is set the default filename is used.
// You can explicitly enable or disable the storage using EnableStoreOf*()
// The default argument is no filename.
//
void MJCut::SetNameSummaryFile(const char *name)
{
    fNameSummary=name;
    if (!fNameSummary.IsNull())
        EnableStorageOfSummary();
}

// --------------------------------------------------------------------------
//
// Set the name of the summary file (events after cut3)
// If you give a name the storage of this file is enabled implicitly.
// If you give no filename the storage is neither enabled nor disabled,
// but the storage file name is reset.
// If no filename is set the default filename is used.
// You can explicitly enable or disable the storage using EnableStoreOf*()
// The default argument is no filename.
//
void MJCut::SetNameResultFile(const char *name)
{
    fNameResult=name;
    if (!fNameResult.IsNull())
        EnableStorageOfResult();
}

// --------------------------------------------------------------------------
//
// Setup a task estimating the energy. The given task is cloned.
//
void MJCut::SetEnergyEstimator(const MTask *task)
{
    if (fEstimateEnergy)
        delete fEstimateEnergy;
    fEstimateEnergy = task ? (MTask*)task->Clone() : 0;
}

// --------------------------------------------------------------------------
//
// Setup a task calculating the hadronness. The given task is cloned.
//
void MJCut::SetHadronnessCalculator(const MTask *task)
{
    if (fCalcHadronness)
        delete fCalcHadronness;
    fCalcHadronness = task ? (MTask*)task->Clone() : 0;
}

// --------------------------------------------------------------------------
//
// return fOutputPath+"/ganymed%08d.root", num
//
TString MJCut::GetOutputFile(UInt_t num) const
{
    TString p(fPathOut);
    p += "/";
    p += fNameOutput.IsNull() ? Form("ganymed%08d.root", num) : fNameOutput.Data();
    return p;
}

/*
Bool_t MJCut::ReadTasks(const char *fname, MTask* &env1, MTask* &env2) const
{
    //    const TString fname = Form("%s/calib%08d.root", fPathIn.Data(), fSequence.GetSequence());

    *fLog << inf << "Reading from file: " << fname << endl;

    TFile file(fname, "READ");
    if (!file.IsOpen())
    {
        *fLog << err << dbginf << "ERROR - Could not open file " << fname << endl;
        return kFALSE;
    }

    TObject *o = file.Get("EstimateEnergy");
    if (o && !o->InheritsFrom(MTask::Class()))
    {
        *fLog << err << dbginf << "ERROR - EstimateEnergy read from " << fname << " doesn't inherit from MTask!" << endl;
        return kFALSE;
    }
    env1 = o ? (MTask*)o->Clone() : NULL;

    o = file.Get("CalcHadronness");
    if (o && !o->InheritsFrom(MTask::Class()))
    {
        *fLog << err << dbginf << "ERROR - CalcHadronness read from " << fname << " doesn't inherit from MTask!" << endl;
        return kFALSE;
    }
    env2 = o ? (MTask*)o->Clone() : NULL;

    return kTRUE;
}
*/

// --------------------------------------------------------------------------
//
// Write the tasks in cont to the file corresponding to analysis number num,
// see GetOutputFile()
//
Bool_t MJCut::WriteTasks(UInt_t num, TObjArray &cont) const
{
    if (fPathOut.IsNull())
    {
        *fLog << inf << "No output path specified via SetPathOut - no output written." << endl;
        return kTRUE;
    }

    const TString oname(GetOutputFile(num));

    *fLog << inf << "Writing to file: " << oname << endl;

    TFile file(oname, fOverwrite?"RECREATE":"NEW", "File created by MJCut", 9);
    if (!file.IsOpen())
    {
        *fLog << err << "ERROR - Couldn't open file " << oname << " for writing..." << endl;
        return kFALSE;
    }

    return WriteContainer(cont);
}

// --------------------------------------------------------------------------
//
// Write the result plots and other results to the file corresponding to
// analysis number num, see GetOutputFile()
//
Bool_t MJCut::WriteResult(UInt_t num) const
{
    TObjArray arr;
    return WriteContainer(arr, GetOutputFile(num), "UPDATE");
}

// --------------------------------------------------------------------------
//
// MJCut allows to setup several option by a resource file:
//   MJCut.WriteSummary: yes, no
//   MJCut.SummaryFile:  filename
//   MJCut.WriteResult:  yes, no
//   MJCut.ResultFile:   filename
//
Bool_t MJCut::CheckEnvLocal()
{
    const TString f0(GetEnv("SummaryFile", ""));
    const TString f1(GetEnv("ResultFile",  ""));
    if (!f0.IsNull())
        SetNameSummaryFile(f0);
    if (!f1.IsNull())
        SetNameResultFile(f1);

    EnableStorageOfSummary(GetEnv("SummaryFile", fStoreSummary));
    EnableStorageOfResult(GetEnv("ResultFile", fStoreResult));
    EnableWobbleMode(GetEnv("WobbleMode", fIsWobble));
    EnableFullDisplay(GetEnv("FullDisplay", fFullDisplay));
    //EnableSubstraction(GetEnv("HistogramSubstraction", fSubstraction));

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Setup write to write:
//     container         tree       optional?
//  --------------     ----------  -----------
//   "MHillas"      to  "Events"
//   "MHillasSrc"   to  "Events"
//   "MHadronness"  to  "Events"       yes
//   "MEnergyEst"   to  "Events"       yes
//   "DataType"     to  "Events"
//
void MJCut::SetupWriter(MWriteRootFile *write, const char *name) const
{
    if (!write)
        return;

    write->SetName(name);
    write->AddContainer("MHillas",        "Events");
    write->AddContainer("MHillasSrc",     "Events");
    write->AddContainer("MPointingPos",   "Events");
    write->AddContainer("MHillasSrcAnti", "Events", kFALSE);
    write->AddContainer("MNewImagePar",   "Events", kFALSE);
    write->AddContainer("MNewImagePar2",  "Events", kFALSE);
    write->AddContainer("MHadronness",    "Events", kFALSE);
    write->AddContainer("MEnergyEst",     "Events", kFALSE);
    write->AddContainer("DataType",       "Events");

    // Should not be the default: Either as option, or as
    // setup from resource file
    // write.AddContainer("MHillasExt",    "Events");
    // write.AddContainer("MImagePar",     "Events");
    // write.AddContainer("MNewImagePar",  "Events");
}

Bool_t MJCut::ProcessFile(const MDataSet &set)
{
    if (!set.IsValid())
    {
        *fLog << err << "ERROR - DataSet invalid!" << endl;
        return kFALSE;
    }

    CheckEnv();

    // --------------------------------------------------------------------------------

    *fLog << inf;
    fLog->Separator(GetDescriptor());
    *fLog << "Perform cuts for data set " << set.GetName() << endl;
    *fLog << endl;

    // --------------------------------------------------------------------------------

    // Setup Parlist
    MParList plist;
    plist.AddToList(this); // take care of fDisplay!

    MParameterI par("DataType");
    plist.AddToList(&par);

    // Setup Tasklist
    MTaskList tlist;
    plist.AddToList(&tlist);

    // La Palma Magic1
    MObservatory obs;
    plist.AddToList(&obs);

    // Possible source position (eg. Wobble Mode)
    MPointingPos source("MSourcePos");
    if (set.GetSourcePos(source))
    {
        plist.AddToList(&source);
        *fLog << inf << "Using Source Position: " << source.GetTitle() << endl;
    }
    else
        *fLog << inf << "No source position applied..." << endl;

    // Initialize default binnings
    MBinning bins1(18, 0,  90,   "BinningAlpha",     "lin");
    MBinning bins2(15, 10, 1e6 , "BinningEnergyEst", "log");
    MBinning bins3(50, 0,  60,   "BinningTheta",     "cos");
    MBinning bins4("BinningFalseSource");
    //MBinning bins5("BinningWidth");
    //MBinning bins6("BinningLength");
    //MBinning bins7("BinningDist");
    plist.AddToList(&bins1);
    plist.AddToList(&bins2);
    plist.AddToList(&bins3);
    plist.AddToList(&bins4);
    //plist.AddToList(&bins5);
    //plist.AddToList(&bins6);
    //plist.AddToList(&bins7);

    // --------------------------------------------------------------------------------

    // Setup fitter and histograms
    MAlphaFitter fit;
    plist.AddToList(&fit);
    if (fIsWobble)
        fit.SetScaleMode(MAlphaFitter::kNone);

    MFillH falpha("MHAlphaOff [MHAlpha]", "MHillasSrc", "FillAlpha");
    MFillH ffs("MHFalseSourceOff [MHFalseSource]", "MHillas", "FillFS");

    // FIXME: If fPathIn read cuts and energy estimator from file!
    MContinue cont0("", "Cut0");
    MContinue cont1("", "Cut1");
    MContinue cont2("", "Cut2");
    MContinue cont3("", "Cut3");
    cont0.SetAllowEmpty();
    cont1.SetAllowEmpty();
    cont2.SetAllowEmpty();
    cont3.SetAllowEmpty();

    // ------------- Loop Off Data --------------------
    MReadReports readoff;
    readoff.AddTree("Events", "MTime.", kTRUE);
    readoff.AddTree("Drive");
    readoff.AddTree("EffectiveOnTime");
    if (fIsWobble)
        set.AddFilesOn(readoff);
    else
        set.AddFilesOff(readoff);

    const TString path(Form("%s/", fPathOut.Data()));
    TString fname0(path);
    TString fname1(path);
    fname0 += fNameSummary.IsNull() ?  (TString) Form("ganymed%08d-summary.root", set.GetNumAnalysis()) : fNameSummary;
    fname1 += fNameResult.IsNull()  ?  (TString) Form("ganymed%08d-result.root",  set.GetNumAnalysis()) : fNameResult;

    MWriteRootFile *write0 = CanStoreSummary() ? new MWriteRootFile(fPathOut.IsNull()?0:fname0.Data(), fOverwrite?"RECREATE":"NEW") : 0;
    MWriteRootFile *write1 = CanStoreResult()  ? new MWriteRootFile(fPathOut.IsNull()?0:fname1.Data(), fOverwrite?"RECREATE":"NEW") : 0;
    SetupWriter(write0, "WriteAfterCut0");
    SetupWriter(write1, "WriteAfterCut3");


    MEnergyEstimate est;

    MTaskEnv taskenv1("EstimateEnergy");
    taskenv1.SetDefault(fEstimateEnergy ? fEstimateEnergy : &est);

    MTaskEnv taskenv2("CalcHadronness");
    taskenv2.SetDefault(fCalcHadronness);

    MFillH fill1a("MHHillasOffPre  [MHHillas]",      "MHillas",      "FillHillasPre");
    MFillH fill2a("MHHillasOffPost [MHHillas]",      "MHillas",      "FillHillasPost");
    MFillH fill3a("MHVsSizeOffPost [MHVsSize]",      "MHillasSrc",   "FillVsSizePost");
    MFillH fill4a("MHHilExtOffPost [MHHillasExt]",   "MHillasSrc",   "FillHilExtPost");
    MFillH fill5a("MHHilSrcOffPost [MHHillasSrc]",   "MHillasSrc",   "FillHilSrcPost");
    MFillH fill6a("MHImgParOffPost [MHImagePar]",    "MImagePar",    "FillImgParPost");
    MFillH fill7a("MHNewParOffPost [MHNewImagePar]", "MNewImagePar", "FillNewParPost");
    fill1a.SetNameTab("PreCut");
    fill2a.SetNameTab("PostCut");
    fill3a.SetNameTab("VsSize");
    fill4a.SetNameTab("HilExt");
    fill5a.SetNameTab("HilSrc");
    fill6a.SetNameTab("ImgPar");
    fill7a.SetNameTab("NewPar");

    MPrint print2("MEffectiveOnTime");

    // How to get source position from off- and on-data?
    MSrcPosCalc scalc;
    if (fIsWobble)
        scalc.SetWobbleMode(); /********************/
    MHillasCalc hcalc;
    MHillasCalc hcalc2("MHillasCalcAnti");
    hcalc.SetFlags(MHillasCalc::kCalcHillasSrc);
    hcalc2.SetFlags(MHillasCalc::kCalcHillasSrc);
    hcalc2.SetNameHillasSrc("MHillasSrcAnti");
    hcalc2.SetNameSrcPosCam("MSrcPosAnti");

    MTaskList tlist2;
    tlist2.AddToList(&scalc);
    tlist2.AddToList(&hcalc);
    if (fIsWobble)
        tlist2.AddToList(&hcalc2);
    tlist2.AddToList(&taskenv1);
    tlist2.AddToList(&taskenv2);
    tlist2.AddToList(&cont0);
    if (write0)
        tlist2.AddToList(write0);
    if (!fWriteOnly)
        tlist2.AddToList(&fill1a);
    tlist2.AddToList(&cont1);
    if (!fWriteOnly && !fIsWobble)
        tlist2.AddToList(&ffs);
    tlist2.AddToList(&cont2);
    if (!fWriteOnly)
    {
        tlist2.AddToList(&fill2a);
        if (fFullDisplay)
        {
            tlist2.AddToList(&fill3a);
            tlist2.AddToList(&fill4a);
            tlist2.AddToList(&fill5a);
            tlist2.AddToList(&fill6a);
            tlist2.AddToList(&fill7a);
        }
    }
    if (!fWriteOnly)
        tlist2.AddToList(&falpha);
    tlist2.AddToList(&cont3);
    if (write1)
        tlist2.AddToList(write1);

    tlist.AddToList(&readoff);
    tlist.AddToList(&print2, "EffectiveOnTime");
    tlist.AddToList(&tlist2, "Events");
 
    par.SetVal(0);

    // Create and setup the eventloop
    MEvtLoop evtloop(fName);
    evtloop.SetParList(&plist);
    evtloop.SetDisplay(fDisplay);
    evtloop.SetLogStream(fLog);
    if (!SetupEnv(evtloop))
        return kFALSE;

    if (set.HasOffSequences() || fIsWobble)
    {
        // Execute first analysis
        if (!evtloop.Eventloop(fMaxEvents))
        {
            *fLog << err << GetDescriptor() << ": Processing of off-sequences failed." << endl;
            return kFALSE;
        }

        tlist.PrintStatistics();

        if (!evtloop.GetDisplay())
        {
            *fLog << err << GetDescriptor() << ": Execution stopped by user." << endl;
            return kFALSE;
        }

        //plist.FindObject("MTimeEffectiveOnTime")->Clear();
    }

    // ------------- Loop On Data --------------------
    MReadReports readon;
    readon.AddTree("Events", "MTime.", kTRUE);
    readon.AddTree("Drive");
    readon.AddTree("EffectiveOnTime");
    set.AddFilesOn(readon);

    if (fIsWobble)
        scalc.SetWobbleMode(kFALSE); /********************/

    MFillH fill1b("MHHillasOnPre  [MHHillas]",      "MHillas",      "FillHillasPre");
    MFillH fill2b("MHHillasOnPost [MHHillas]",      "MHillas",      "FillHillasPost");
    MFillH fill3b("MHVsSizeOnPost [MHVsSize]",      "MHillasSrc",   "FillVsSizePost");
    MFillH fill4b("MHHilExtOnPost [MHHillasExt]",   "MHillasSrc",   "FillHilExtPost");
    MFillH fill5b("MHHilSrcOnPost [MHHillasSrc]",   "MHillasSrc",   "FillHilSrcPost");
    MFillH fill6b("MHImgParOnPost [MHImagePar]",    "MImagePar",    "FillImgParPost");
    MFillH fill7b("MHNewParOnPost [MHNewImagePar]", "MNewImagePar", "FillNewParPost");
    fill1b.SetNameTab("PreCut");
    fill2b.SetNameTab("PostCut");
    fill3b.SetNameTab("VsSize");
    fill4b.SetNameTab("HilExt");
    fill5b.SetNameTab("HilSrc");
    fill6b.SetNameTab("ImgPar");
    fill7b.SetNameTab("NewPar");
    fill1b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill2b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill3b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill4b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill5b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill6b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill7b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");

    /*
     MHVsTime hvs("MEffectiveOnTime.fVal");
     hvs.SetTitle("Effective On-Time vs. Time;;T_{on}");
     MFillH fillvs(&hvs, "MTimeEffectiveOnTime", "FillOnTime");
     fillvs.SetNameTab("OnTime");
     */
    MH3 hvs("MPointingPos.fZd");
    hvs.SetName("Theta");
    hvs.SetTitle("Effective On-Time vs. Zenith Angle;\\Theta [\\circ];T_{on} [s]");
    MFillH fillvs(&hvs, "", "FillOnTime");
    fillvs.SetWeight("MEffectiveOnTime");
    fillvs.SetNameTab("OnTime");

    /*
    MParameterD weight;
    weight.SetVal(-1);
    fill2a.SetWeight(&weight);
    fill3a.SetWeight(&weight);
    fill4a.SetWeight(&weight);
    fill5a.SetWeight(&weight);
    fill6a.SetWeight(&weight);
    fill7a.SetWeight(&weight);
    if (fSubstraction)
    {
        fill2a.SetNameTab("PostCut-");
        fill3a.SetNameTab("VsSize-");
        fill4a.SetNameTab("HilExt-");
        fill5a.SetNameTab("HilSrc-");
        fill6a.SetNameTab("ImgPar-");
        fill7a.SetNameTab("NewPar-");
    }
    */
    MFillH falpha2("MHAlpha", "MHillasSrc", "FillAlpha");
    MFillH ffs2("MHFalseSource", "MHillas", "FillFS");

    tlist.Replace(&readon);
    if (!fWriteOnly)
    {
        tlist2.Replace(&fill1b);
/*        if (fIsWobble)
        {
            tlist2.AddToListAfter(&fill2b, &fill2a);
            tlist2.AddToListAfter(&fill3b, &fill3a);
        }
        else
        */
        tlist2.Replace(&fill2b);
        if (fFullDisplay)
        {
            tlist2.Replace(&fill3b);
            tlist2.Replace(&fill4b);
            tlist2.Replace(&fill5b);
            tlist2.Replace(&fill6b);
            tlist2.Replace(&fill7b);
        }
        tlist2.Replace(&falpha2);
        if (!fIsWobble)
            tlist2.Replace(&ffs2);

        tlist.AddToList(&fillvs, "EffectiveOnTime");
    }

    par.SetVal(1);

    TObjArray cont;
    cont.Add(&cont0);
    cont.Add(&cont1);
    cont.Add(&cont2);
    cont.Add(&cont3);
    if (taskenv1.GetTask())
        cont.Add(taskenv1.GetTask());
    if (taskenv2.GetTask())
        cont.Add(taskenv2.GetTask());

    if (!WriteTasks(set.GetNumAnalysis(), cont))
        return kFALSE;

    // Execute first analysis
    if (!evtloop.Eventloop(fMaxEvents))
    {
        *fLog << err << GetDescriptor() << ": Processing of on-sequences failed." << endl;
        return kFALSE;
    }

    if (write0)
        delete write0;
    if (write1)
        delete write1;

    tlist.PrintStatistics();

    // FIXME: Perform fit and plot energy dependant alpha plots
    // and fit result to new tabs!
    if (!WriteResult(set.GetNumAnalysis()))
        return kFALSE;

    *fLog << all << GetDescriptor() << ": Done." << endl;
    *fLog << endl << endl;

    return kTRUE;
}
