/* ======================================================================== *\
!
! *
! * 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 "MBinning.h"
#include "MDataSet.h"
#include "MParameters.h"
#include "MPointingPos.h"
#include "MObservatory.h"

ClassImp(MJCut);

using namespace std;

// --------------------------------------------------------------------------
//
// Default constructor. 
//
MJCut::MJCut(const char *name, const char *title)
    : fStoreSummary(kFALSE), fStoreResult(kFALSE), fWriteOnly(kFALSE),
    fIsWobble(kFALSE), fEstimateEnergy(0), fCalcHadronness(0)
{
    fName  = name  ? name  : "MJCut";
    fTitle = title ? title : "Standard program to perform g/h-seperation cuts";
}

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));

    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
{
    write.SetName(name);
    write.AddContainer("MHillas",        "Events");
    write.AddContainer("MHillasSrc",     "Events");
    write.AddContainer("MHillasSrcAnti", "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");
    plist.AddToList(&bins1);
    plist.AddToList(&bins2);
    plist.AddToList(&bins3);
    plist.AddToList(&bins4);

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

    // 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(fPathOut.IsNull()?0:fname0.Data(), fOverwrite?"RECREATE":"NEW");
    MWriteRootFile write1(fPathOut.IsNull()?0:fname1.Data(), fOverwrite?"RECREATE":"NEW");
    if (CanStoreSummary())
        SetupWriter(write0, "WriteAfterCut0");
    if (CanStoreSummary())
        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");
    fill1a.SetNameTab("PreCut");
    fill2a.SetNameTab("PostCut");
    fill3a.SetNameTab("VsSize");

    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 (CanStoreSummary())
        tlist2.AddToList(&write0);
    if (!fWriteOnly)
        tlist2.AddToList(&fill1a);
    tlist2.AddToList(&cont1);
    if (!fWriteOnly)
        tlist2.AddToList(&ffs);
    tlist2.AddToList(&cont2);
    if (!fWriteOnly)
    {
        tlist2.AddToList(&fill2a);
        tlist2.AddToList(&fill3a);
    }
    if (!fWriteOnly)
        tlist2.AddToList(&falpha);
    tlist2.AddToList(&cont3);
    if (CanStoreResult())
        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;
        }
    }

    // ------------- 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");
    fill1b.SetNameTab("PreCut");
    fill2b.SetNameTab("PostCut");
    fill3b.SetNameTab("VsSize");
    fill1b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill2b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");
    fill3b.SetDrawOption(set.HasOffSequences()||fIsWobble?"same":"");

    MFillH falpha2("MHAlpha", "MHillasSrc", "FillAlpha");
    MFillH ffs2("MHFalseSource", "MHillas", "FillFS");

    tlist.Replace(&readon);
    if (!fWriteOnly)
    {
        tlist2.Replace(&fill1b);
        tlist2.Replace(&fill2b);
        tlist2.Replace(&fill3b);
        tlist2.Replace(&falpha2);
        tlist2.Replace(&ffs2);
    }

    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;
    }

    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;
}