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

/////////////////////////////////////////////////////////////////////////////
//
//  MJSpectrum
//
// Program to calculate spectrum
//
/////////////////////////////////////////////////////////////////////////////
#include "MJSpectrum.h"

// Root
#include <TF1.h>
#include <TH1.h>
#include <TH2.h>
#include <TFile.h>
#include <TChain.h>
#include <TCanvas.h>
#include <TObjArray.h>

// Environment
#include "MLog.h"
#include "MLogManip.h"

#include "MStatusArray.h"
#include "MStatusDisplay.h"

// Container
#include "MH3.h"
#include "MBinning.h"
#include "MDataSet.h"

// Spectrum
#include "../mhflux/MAlphaFitter.h"
#include "../mhflux/MHAlpha.h"
#include "../mhflux/MHCollectionArea.h"
#include "../mhflux/MHEnergyEst.h"

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

// Tasks/Filter
#include "MReadMarsFile.h"
#include "MFEventSelector2.h"
#include "MFDataMember.h"
#include "MFillH.h"
#include "MHillasCalc.h"
#include "MSrcPosCalc.h"
#include "MContinue.h"

ClassImp(MJSpectrum);

using namespace std;

MJSpectrum::MJSpectrum(const char *name, const char *title)
    : fCut0(0),fCut1(0), fCut2(0), fCut3(0), fEstimateEnergy(0),
    fRefill(kFALSE), fSimpleMode(kTRUE)
{
    fName  = name  ? name  : "MJSpectrum";
    fTitle = title ? title : "Standard program to calculate spectrum";
}

MJSpectrum::~MJSpectrum()
{
    if (fCut0)
        delete fCut0;
    if (fCut1)
        delete fCut1;
    if (fCut2)
        delete fCut2;
    if (fCut3)
        delete fCut3;
    if (fEstimateEnergy)
        delete fEstimateEnergy;
}

Bool_t MJSpectrum::ReadTask(MTask* &task, const char *name) const
{
    if (task)
    {
        delete task;
        task = 0;
    }

    task = (MTask*)gFile->Get(name);
    if (!task)
    {
        *fLog << err << dbginf << "ERROR - " << name << " doen't exist in file!" << endl;
        return kFALSE;
    }
    if (!task->InheritsFrom(MTask::Class()))
    {
        *fLog << err << dbginf << "ERROR - " << name << " read doesn't inherit from MTask!" << endl;
        delete task;
        return kFALSE;
    }

    task->SetName(name);

    return kTRUE;
}

Bool_t MJSpectrum::ReadInput(const MParList &plist)
{
    const TString fname = fPathIn;

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

    if (!ReadTask(fCut0, "Cut0"))
        return kFALSE;
    if (!ReadTask(fCut1, "Cut1"))
        return kFALSE;
    if (!ReadTask(fCut2, "Cut2"))
        return kFALSE;
    if (!ReadTask(fCut3, "Cut3"))
        return kFALSE;

    if (!ReadTask(fEstimateEnergy, "EstimateEnergy"))
        return kFALSE;

    TObjArray arr;

    TIter Next(plist);
    TObject *o=0;
    while ((o=Next()))
        if (o->InheritsFrom(MBinning::Class()))
            arr.Add(o);

    arr.Add(plist.FindObject("MAlphaFitter"));

    return ReadContainer(arr);
}

void MJSpectrum::PrintSetup(const MAlphaFitter &fit) const
{
    gLog.Separator("Alpha Fitter");
    *fLog << all;
    fit.Print();

    gLog.Separator("Used Cuts");
    fCut0->Print();
    fCut1->Print();
    fCut2->Print();
    fCut3->Print();

    gLog.Separator("Energy Estimator");
    fEstimateEnergy->Print();
}

Float_t MJSpectrum::ReadHistograms(TH1D &h1, TH1D &h2) const
{
    TFile file(fPathIn, "READ");
    if (!file.IsOpen())
    {
        *fLog << err << dbginf << "ERROR - Could not open file " << fPathIn << endl;
        return -1;
    }

    MStatusArray *arr = (MStatusArray*)file.Get("MStatusDisplay");
    if (!arr)
    {
        gLog << "MStatusDisplay not found in file... abort." << endl;
        return -1;
    }

    TH1D *vstime = (TH1D*)arr->FindObjectInCanvas("Theta",        "TH1D", "OnTime");
    TH1D *excen  = (TH1D*)arr->FindObjectInCanvas("ExcessEnergy", "TH1D", "MHAlpha");
    if (!vstime || !excen)
        return -1;

    vstime->Copy(h1);
    excen->Copy(h2);
    h1.SetDirectory(0);
    h2.SetDirectory(0);

    if (fDisplay)
        arr->DisplayIn(*fDisplay, "MHAlpha");

    delete arr;

    return vstime->Integral();
}

Bool_t MJSpectrum::ReadOrigMCDistribution(const MDataSet &set, TH1 &h) const
{
    // Some debug output
    fLog->Separator("Compiling original MC distribution");

    *fLog << inf << "Please stand by, this may take a while..." << flush;

    if (fDisplay)
        fDisplay->SetStatusLine1("Compiling MC distribution...");

    // Create chain
    TChain chain("OriginalMC");
    set.AddFilesOn(chain);

    // Prepare histogram
    h.Reset();

    // Fill histogram from chain
    h.SetDirectory(gROOT);
    if (h.InheritsFrom(TH2::Class()))
    {
        h.SetNameTitle("ThetaEMC", "Event-Distribution vs Theta and Energy for MC (produced)");
        h.SetXTitle("\\Theta [\\circ]");
        h.SetYTitle("E [GeV]");
        h.SetZTitle("Counts");
        chain.Draw("MMcEvtBasic.fEnergy:MMcEvtBasic.fTelescopeTheta*TMath::RadToDeg()>>ThetaEMC", "", "goff");
    }
    else
    {
        h.SetNameTitle("ThetaMC", "Event-Distribution vs Theta for MC (produced)");
        h.SetXTitle("\\Theta [\\circ]");
        h.SetYTitle("Counts");
        chain.Draw("MMcEvtBasic.fTelescopeTheta*TMath::RadToDeg()>>ThetaMC", "", "goff");
    }
    h.SetDirectory(0);

    *fLog << "done." << endl;
    if (fDisplay)
        fDisplay->SetStatusLine2("done.");

    if (h.GetEntries()>0)
        return kTRUE;

    *fLog << err << "ERROR - Histogram with original MC distribution empty..." << endl;

    return h.GetEntries()>0;
}

Bool_t MJSpectrum::GetThetaDistribution(TH1D &temp1, TH1D &temp2) const
{
    // Display some stuff
    if (fDisplay)
    {
        TCanvas &c = fDisplay->AddTab("ZdDist");
        c.Divide(2,2);

        // On-Time vs. Theta
        c.cd(1);
        gPad->SetBorderMode(0);
        temp1.DrawCopy();

        // Number of MC events (produced) vs Theta
        c.cd(2);
        gPad->SetBorderMode(0);
        temp2.SetName("NVsTheta");
        temp2.DrawCopy();

        c.cd(4);
        gPad->SetBorderMode(0);

        c.cd(3);
        gPad->SetBorderMode(0);
    }

    // Calculate the Probability
    temp1.Divide(&temp2);
    temp1.Scale(1./temp1.GetMaximum());

    // Some cosmetics: Name, Axis, etc.
    temp1.SetName("ProbVsTheta");
    temp1.SetTitle("Probability vs. Zenith Angle to choose MC events");
    temp1.SetYTitle("Probability");
    if (fDisplay)
        temp1.DrawCopy();

    return kTRUE;
}

void MJSpectrum::DisplayResult(const TH2D &h2) const
{
    if (!fDisplay->CdCanvas("ZdDist"))
        return;

    TH1D &proj = *h2.ProjectionX();
    proj.SetNameTitle("ThetaFinal", "Final Theta Distribution");
    proj.SetXTitle("\\Theta [\\circ]");
    proj.SetYTitle("Counts");
    proj.SetLineColor(kBlue);
    proj.SetDirectory(0);
    proj.SetBit(kCanDelete);

    TVirtualPad *pad = gPad;

    pad->cd(4);
    proj.DrawCopy();

    pad->cd(1);
    TH1D *theta = (TH1D*)gPad->FindObject("Theta");
    if (theta)
    {
        proj.Scale(theta->Integral()/proj.Integral());
        theta->SetMaximum(1.05*TMath::Max(theta->GetMaximum(), proj.GetMaximum()));
    }
    proj.Draw("same");
}

Bool_t MJSpectrum::Refill(MParList &plist, TH1D &h2) const
{
    *fLog << endl;
    fLog->Separator("Refill Excess");
    *fLog << endl;

    MTaskList tlist;
    plist.AddToList(&tlist);

    MReadTree read("Events");
    read.DisableAutoScheme();
    read.AddFile(fPathIn);

    MFillH fill1("MHAlphaOff [MHAlpha]", "MHillasSrc", "FillAlpha");
    MFillH fill2("MHAlpha", "MHillasSrc", "FillAlpha");

    MFDataMember f0("DataType.fVal", '<', 0.5, "FilterOffData");
    MFDataMember f1("DataType.fVal", '>', 0.5, "FilterOnData");

    fill1.SetFilter(&f0);
    fill2.SetFilter(&f1);

    tlist.AddToList(&read);
    tlist.AddToList(fEstimateEnergy);
    tlist.AddToList(&f0);
    tlist.AddToList(&f1);
    tlist.AddToList(&fill1);
    tlist.AddToList(&fill2);

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

    if (!SetupEnv(loop))
        return kFALSE;

    if (!loop.Eventloop())
    {
        *fLog << err << GetDescriptor() << ": Refilling of data failed." << endl;
        return kFALSE;
    }

    tlist.PrintStatistics();

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

    const MHAlpha *halpha = (MHAlpha *)plist.FindObject("MHAlpha");
    if (!halpha)
    {
        *fLog << err << GetDescriptor() << ": MHAlpha not found... abort." << endl;
        return kFALSE;
    }

    halpha->GetHEnergy().Copy(h2);
    h2.SetDirectory(0);

    return kTRUE;
}

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

    CheckEnv();

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

    *fLog << inf;
    fLog->Separator(GetDescriptor());
    *fLog << "Compile Monte Carlo Sample (data set " << set.GetName() << ")" << endl;
    *fLog << endl;

    MBinning bins1("BinningAlpha");
    MBinning bins2("BinningEnergyEst");
    MBinning bins3("BinningTheta");
    MBinning bins4("BinningFalseSource");
    MBinning bins5("BinningWidth");
    MBinning bins6("BinningLength");
    MBinning bins7("BinningDist");
    MBinning bins8("BinningMaxDist");

    MAlphaFitter fit;

    MParList plist;
    plist.AddToList(&bins1);
    plist.AddToList(&bins2);
    plist.AddToList(&bins3);
    plist.AddToList(&bins4);
    plist.AddToList(&bins5);
    plist.AddToList(&bins6);
    plist.AddToList(&bins7);
    plist.AddToList(&bins8);
    plist.AddToList(&fit);

    if (!ReadInput(plist))
        return kFALSE;

    PrintSetup(fit);

    TH1D temp1, excess;
    const Float_t ontime = ReadHistograms(temp1, excess);
    if (ontime<0)
    {
        *fLog << err << GetDescriptor() << ": Could not determin effective on time..." << endl;
        return kFALSE;
    }

    TH1D temp2(temp1);
    if (!ReadOrigMCDistribution(set, temp2))
        return kFALSE;

    if (!GetThetaDistribution(temp1, temp2))
        return kFALSE;

    if (fRefill && !Refill(plist, temp2))
        return kFALSE;

    TH2D hist;
    MH3 mh1("MMcEvtBasic.fTelescopeTheta*kRad2Deg", "MMcEvtBasic.fEnergy");
    if (fSimpleMode)
    {
        hist.UseCurrentStyle();
        MH::SetBinning(&hist, temp1.GetXaxis(), excess.GetXaxis());
        if (!ReadOrigMCDistribution(set, hist))
            return kFALSE;

        for (int y=0; y<hist.GetNbinsY(); y++)
            for (int x=0; x<hist.GetNbinsX(); x++)
                hist.SetBinContent(x, y, hist.GetBinContent(x, y)*temp1.GetBinContent(x));
    }
    else
    {
        MTaskList tlist1;
        plist.Replace(&tlist1);

        MReadMarsFile readmc("OriginalMC");
        //readmc.DisableAutoScheme();
        set.AddFilesOn(readmc);
        readmc.EnableBranch("MMcEvtBasic.fTelescopeTheta");
        readmc.EnableBranch("MMcEvtBasic.fEnergy");

        mh1.SetLogy();
        mh1.SetLogz();
        mh1.SetName("ThetaE");

        MFillH fill0(&mh1);
        //fill0.SetDrawOption("projx only");

        MBinning binsx(bins3, "BinningThetaEX");
        MBinning binsy(bins2, "BinningThetaEY");
        plist.AddToList(&binsx);
        plist.AddToList(&binsy);
        tlist1.AddToList(&readmc);

        temp1.SetXTitle("MMcEvtBasic.fTelescopeTheta*kRad2Deg");
        MH3 mh3mc(temp1);

        MFEventSelector2 sel1(mh3mc);
        sel1.SetHistIsProbability();

        fill0.SetFilter(&sel1);

        tlist1.AddToList(&sel1);
        tlist1.AddToList(&fill0);

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

        if (!SetupEnv(loop1))
            return kFALSE;

        if (!loop1.Eventloop(fMaxEvents))
        {
            *fLog << err << GetDescriptor() << ": Processing of MC-data failed." << endl;
            return kFALSE;
        }

        tlist1.PrintStatistics();

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

    DisplayResult(fSimpleMode ? hist : (TH2D&)mh1.GetHist());

    // ------------------------- Final loop --------------------------

    *fLog << endl;
    fLog->Separator("Calculate efficiencies");
    *fLog << endl;

    MTaskList tlist2;
    plist.Replace(&tlist2);

    MReadMarsFile read("Events");
    read.DisableAutoScheme();
    set.AddFilesOn(read);

    // Selector to get correct (final) theta-distribution
    temp1.SetXTitle("MPointingPos.fZd");

    MH3 mh3(temp1);

    MFEventSelector2 sel2(mh3);
    sel2.SetHistIsProbability();

    MContinue contsel(&sel2);
    contsel.SetInverted();

    // Get correct source position
    MSrcPosCalc calc;

    // Calculate corresponding Hillas parameters
    MHillasCalc hcalc1;
    MHillasCalc hcalc2("MHillasCalcAnti");
    hcalc1.SetFlags(MHillasCalc::kCalcHillasSrc);
    hcalc2.SetFlags(MHillasCalc::kCalcHillasSrc);
    hcalc2.SetNameHillasSrc("MHillasSrcAnti");
    hcalc2.SetNameSrcPosCam("MSrcPosAnti");

    // Fill collection area and energy estimator (unfolding)
    // Make sure to use the same binning for MHCollectionArea and MHEnergyEst
    MHCollectionArea area;
    area.SetHistAll(fSimpleMode ? hist : (TH2D&)mh1.GetHist());
    MHEnergyEst      hest;

    MFillH fill3(&area, "", "FillCollectionArea");
    MFillH fill4(&hest, "", "FillEnergyEst");

    MFillH fill1a("MHHillasMCPre  [MHHillas]",      "MHillas",      "FillHillasPre");
    MFillH fill2a("MHHillasMCPost [MHHillas]",      "MHillas",      "FillHillasPost");
    MFillH fill3a("MHVsSizeMCPost [MHVsSize]",      "MHillasSrc",   "FillVsSizePost");
    MFillH fill4a("MHHilExtMCPost [MHHillasExt]",   "MHillasSrc",   "FillHilExtPost");
    MFillH fill5a("MHHilSrcMCPost [MHHillasSrc]",   "MHillasSrc",   "FillHilSrcPost");
    MFillH fill6a("MHImgParMCPost [MHImagePar]",    "MImagePar",    "FillImgParPost");
    MFillH fill7a("MHNewParMCPost [MHNewImagePar]", "MNewImagePar", "FillNewParPost");
    fill1a.SetNameTab("PreCut");
    fill2a.SetNameTab("PostCut");
    fill3a.SetNameTab("VsSize");
    fill4a.SetNameTab("HilExt");
    fill5a.SetNameTab("HilSrc");
    fill6a.SetNameTab("ImgPar");
    fill7a.SetNameTab("NewPar");

    tlist2.AddToList(&read);
    tlist2.AddToList(&contsel);
    tlist2.AddToList(&calc);
    tlist2.AddToList(&hcalc1);
    tlist2.AddToList(&hcalc2);
    tlist2.AddToList(&fill1a);
    tlist2.AddToList(fCut0);
    tlist2.AddToList(fCut1);
    tlist2.AddToList(fCut2);
    tlist2.AddToList(fCut3);
    tlist2.AddToList(fEstimateEnergy);
    tlist2.AddToList(&fill3);
    tlist2.AddToList(&fill4);
    tlist2.AddToList(&fill2a);
    tlist2.AddToList(&fill3a);
    tlist2.AddToList(&fill4a);
    tlist2.AddToList(&fill5a);
    tlist2.AddToList(&fill6a);
    tlist2.AddToList(&fill7a);

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

    if (!SetupEnv(loop2))
        return kFALSE;

    if (!loop2.Eventloop(fMaxEvents))
    {
        *fLog << err << GetDescriptor() << ": Processing of MC-data failed." << endl;
        return kFALSE;
    }

    tlist2.PrintStatistics();

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

    // -------------------------- Spectrum ----------------------------

    TH1D collarea(area.GetHEnergy());
    TH1D weights;
    hest.GetWeights(weights);

    cout << "Effective On time: " << ontime << "s" << endl;

    excess.SetDirectory(NULL);
    excess.SetBit(kCanDelete);
    excess.Scale(1./ontime);
    excess.Divide(&collarea);
    excess.SetNameTitle("Preliminary", "N/sm^{2} versus Energy (before unfolding)");
    excess.SetYTitle("N/sm^{2}");

    TCanvas &c = fDisplay->AddTab("Spectrum");
    c.Divide(2,2);
    c.cd(1);
    gPad->SetBorderMode(0);
    gPad->SetLogx();
    gPad->SetLogy();
    gPad->SetGridx();
    gPad->SetGridy();
    collarea.DrawCopy();

    c.cd(2);
    gPad->SetBorderMode(0);
    gPad->SetLogx();
    gPad->SetLogy();
    gPad->SetGridx();
    gPad->SetGridy();
    excess.DrawCopy();

    c.cd(3);
    gPad->SetBorderMode(0);
    gPad->SetLogx();
    gPad->SetLogy();
    gPad->SetGridx();
    gPad->SetGridy();
    weights.DrawCopy();

    excess.Divide(&weights);
    //excess.Multiply(&weights);
    excess.SetNameTitle("Flux", "N/TeVsm^{2} versus Energy (after unfolding)");

    for (int i=0; i<excess.GetNbinsX(); i++)
    {
        excess.SetBinContent(i+1, excess.GetBinContent(i+1)/excess.GetBinWidth(i+1)*1000);
        excess.SetBinError(i+1,   excess.GetBinError(i+1)/  excess.GetBinWidth(i+1)*1000);
    }

    excess.SetYTitle("N/TeVsm^{2}");

    c.cd(4);
    gPad->SetBorderMode(0);
    gPad->SetLogx();
    gPad->SetLogy();
    gPad->SetGridx();
    gPad->SetGridy();
    excess.DrawCopy();

    TF1 f("f", "[1]*(x/1e3)^[0]", 50, 3e4);
    f.SetParameter(0, -2.87);
    f.SetParameter(1, 1.9e-6);
    f.SetLineColor(kGreen);
    excess.Fit(&f, "NI", "", 55, 2e4);
    f.DrawCopy("same");

    if (!fPathOut.IsNull())
        fDisplay->SaveAsRoot(fPathOut);

    /*
     TString str;
     str += "(1.68#pm0.15)10^{-7}";
     str += "(\\frac{E}{TeV})^{-2.59#pm0.06}";
     str += "\\frac{ph}{TeVm^{2}s}";

     TLatex tex;
     tex.DrawLatex(2e2, 7e-5, str);
     */

    return kTRUE;
}