/* ======================================================================== *\
!
! *
! * 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): Wolfgang Wittek 4/2002 <mailto:wittek@mppmu.mpg.de>
!   Author(s): Abelardo Moralejo 5/2003 <mailto:moralejo@pd.infn.it>
!   Author(s): Thomas Bretz 1/2005 <mailto:tbretz@astro.uni-wuerzburg.de>
!
!   Copyright: MAGIC Software Development, 2000-2005
!
!
\* ======================================================================== */

//////////////////////////////////////////////////////////////////////////////
//
//  MHEnergyEst
//
//  calculates the migration matrix E-est vs. E-true
//  for different bins in Theta
//
//////////////////////////////////////////////////////////////////////////////
#include "MHEnergyEst.h"

#include <TF1.h>
#include <TLine.h>
#include <TCanvas.h>
#include <TStyle.h>
#include <TProfile.h>
#include <TPaveStats.h>

#include "MString.h"

#include "MMcEvt.hxx"

#include "MBinning.h"
#include "MParList.h"
#include "MParameters.h"
#include "MHMatrix.h"

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

ClassImp(MHEnergyEst);

using namespace std;

// --------------------------------------------------------------------------
//
// Default Constructor. It sets name and title of the histogram.
//
MHEnergyEst::MHEnergyEst(const char *name, const char *title)
    : fMcEvt(0), fEnergy(0), fResult(0), fMatrix(0)
{
    //
    //   set the name and title of this object
    //
    fName  = name  ? name  : "MHEnergyEst";
    fTitle = title ? title : "Histogram for the result of the energy reconstruction";

    //fNameEnergy = "MEnergyEst";
    //fNameResult = "MinimizationValue";

    fHEnergy.SetDirectory(NULL);
    fHEnergy.SetName("EnergyEst");
    fHEnergy.SetTitle("Histogram in E_{est}, E_{mc} and \\Theta");
    fHEnergy.SetXTitle("E_{est} [GeV]");
    fHEnergy.SetYTitle("E_{mc} [GeV]");
    fHEnergy.SetZTitle("\\Theta [\\circ]");

    fHResolution.SetDirectory(NULL);
    fHResolution.SetName("EnergyRes");
    fHResolution.SetTitle("Histogram in \\Delta E/E vs E_{est} and E_{mc}");
    fHResolution.SetXTitle("E_{est} [GeV]");
    fHResolution.SetYTitle("E_{mc} [GeV]");
    fHResolution.SetZTitle("E_{est}/E_{mc}-1");

    fHImpact.SetDirectory(NULL);
    fHImpact.SetName("Impact");
    fHImpact.SetTitle("\\Delta E/E vs Impact parameter");
    fHImpact.SetXTitle("Impact parameter [m]");
    fHImpact.SetYTitle("E_{est}/E_{mc}-1");

    MBinning binsi, binse, binst, binsr;
    binse.SetEdgesLog(25, 10, 1000000);
    binst.SetEdgesASin(51, -0.005, 0.505);
    binsi.SetEdges(10, 0, 400);
    binsr.SetEdges(75, -1.75, 1.75);

    SetBinning(&fHEnergy,     &binse, &binse, &binst);
    SetBinning(&fHImpact,     &binsi, &binsr);
    SetBinning(&fHResolution, &binse, &binse, &binsr);

    // For some unknown reasons this must be called after
    // the binning has been initialized at least once
    fHEnergy.Sumw2();
    fHResolution.Sumw2();
    fHImpact.Sumw2();
}

// --------------------------------------------------------------------------
//
// Set the binnings and prepare the filling of the histograms
//
Bool_t MHEnergyEst::SetupFill(const MParList *plist)
{
    if (!fMatrix)
    {
        fMcEvt = (MMcEvt*)plist->FindObject("MMcEvt");
        if (!fMcEvt)
        {
            *fLog << err << "MMcEvt not found... aborting." << endl;
            return kFALSE;
        }
    }

    fEnergy = (MParameterD*)plist->FindObject("MEnergyEst", "MParameterD");
    if (!fEnergy)
    {
        *fLog << err << "MEnergyEst [MParameterD] not found... aborting." << endl;
        return kFALSE;
    }

    fResult = (MParameterD*)const_cast<MParList*>(plist)->FindCreateObj("MParameterD", "MinimizationValue");
    if (!fResult)
        return kFALSE;

    MBinning binst, binse, binsi, binsr;
    binse.SetEdges(fHEnergy, 'x');
    binst.SetEdges(fHEnergy, 'z');
    binsi.SetEdges(fHImpact, 'x');
    binsr.SetEdges(fHImpact, 'y');

    binst.SetEdges(*plist, "BinningTheta");
    binse.SetEdges(*plist, "BinningEnergyEst");
    binsi.SetEdges(*plist, "BinningImpact");
    binsr.SetEdges(*plist, "BinningEnergyRes");

    SetBinning(&fHEnergy,     &binse, &binse, &binst);
    SetBinning(&fHImpact,     &binsi, &binsr);
    SetBinning(&fHResolution, &binse, &binse, &binsr);

    fChisq = 0;
    fBias  = 0;

    fHEnergy.Reset();
    fHImpact.Reset();
    fHResolution.Reset();

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Fill the histogram
//
Bool_t MHEnergyEst::Fill(const MParContainer *par, const Stat_t w)
{
    const Double_t eest  = fEnergy->GetVal();
    const Double_t etru  = fMatrix ? GetVal(0) : fMcEvt->GetEnergy();
    const Double_t imp   = fMatrix ? GetVal(1) : fMcEvt->GetImpact()/100;
    const Double_t theta = fMatrix ? GetVal(2) : fMcEvt->GetTelescopeTheta()*TMath::RadToDeg();
    const Double_t resE  = (eest-etru)/etru;

    fHEnergy.Fill(eest, etru, theta, w);
    fHResolution.Fill(eest, etru, resE, w);
    fHImpact.Fill(imp, resE, w);

    // For the fit we use a different quantity
    //const Double_t res = TMath::Log10(eest/etru);
    const Double_t res = eest-etru;

    fChisq += res*res;
    fBias  += res;

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Divide chisq and bias by number of executions
// Print result
//
Bool_t MHEnergyEst::Finalize()
{
    fChisq /= GetNumExecutions();
    fBias  /= GetNumExecutions();

    fResult->SetVal(fChisq);

    *fLog << all << endl;
    Print();

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Print result
//
void MHEnergyEst::Print(Option_t *o) const
{
    const Double_t res = TMath::Sqrt(fChisq-fBias*fBias);
    if (!TMath::Finite(res))
    {
        *fLog << all << "MHEnergyEst::Print: ERROR - Resolution is not finite (eg. NaN)." << endl;
        return;
    }

    TH1D *h = (TH1D*)fHResolution.ProjectionZ("Resolution");
    h->Fit("gaus", "Q0", "", -1.0, 0.25);

    TF1 *f = h->GetFunction("gaus");

    *fLog << all << "F=" << fChisq << endl;
    *fLog << "Results from Histogram:" << endl;
    *fLog << " Mean  of Delta E/E: " << Form("%+4.2f%%", 100*h->GetMean()) << endl;
    *fLog << " RMS   of Delta E/E: " << Form("%4.2f%%",  100*h->GetRMS()) << endl;
    *fLog << "Results from Histogram-Fit:" << endl;
    *fLog << " Bias  of Delta E/E: " << Form("%+4.2f%%", 100*f->GetParameter(1)) << endl;
    *fLog << " Sigma of Delta E/E: " << Form("%4.2f%%",  100*f->GetParameter(2)) << endl;
    *fLog << " Res   of Delta E/E: " << Form("%4.2f%%",  100*TMath::Hypot(f->GetParameter(1), f->GetParameter(2))) << endl;

    delete h;
}

// --------------------------------------------------------------------------
//
// Return Correction Coefficients (weights)
// hist = E_mc/E_est
//
void MHEnergyEst::GetWeights(TH1D &hist) const
{
    // Project into EnergyEst_ey
    // the "e" ensures that errors are calculated
    TH1D *h1 = (TH1D*)fHEnergy.Project3D("rtlprmft_ex"); // E_Est
    TH1D *h2 = (TH1D*)fHEnergy.Project3D("rtlprmft_ey"); // E_mc

    h2->Copy(hist);
    hist.Divide(h1);

    delete h1;
    delete h2;

    hist.SetNameTitle("EnergyRatio", "Ratio between estimated and monte carlo energy");
    hist.SetXTitle("E [GeV]");
    hist.SetYTitle("N_{mc}/N_{est} [1]");
    hist.SetDirectory(0);
}

void MHEnergyEst::Paint(Option_t *opt)
{
    TVirtualPad *pad = gPad;

    pad->cd(1);

    TH1D *hx=0;
    TH1D *hy=0;

    if (pad->GetPad(1))
    {
        pad->GetPad(1)->cd(1);

        if ((hx=(TH1D*)gPad->FindObject("EnergyEst_ex")))
        {
            TH1D *h2 = (TH1D*)fHEnergy.Project3D("dum_ex");
            hx->Reset();
            hx->Add(h2);
            delete h2;
        }

        if ((hy=(TH1D*)gPad->FindObject("EnergyEst_ey")))
        {
            TH1D *h2 = (TH1D*)fHEnergy.Project3D("dum_ey");
            hy->Reset();
            hy->Add(h2);
            delete h2;
        }

        if (hx && hy)
        {
            hy->SetMaximum();
            hy->SetMaximum(TMath::Max(hx->GetMaximum(), hy->GetMaximum())*1.2);
            if (hy->GetMaximum()>0)
                gPad->SetLogy();
        }

        if (pad->GetPad(1)->GetPad(2))
        {
            pad->GetPad(1)->GetPad(2)->cd(1);
            if ((hx=(TH1D*)gPad->FindObject("EnergyEst_ez")))
            {
                TH1D *h2 = (TH1D*)fHEnergy.Project3D("dum_ez");
                hx->Reset();
                hx->Add(h2);
                delete h2;
            }

            pad->GetPad(1)->GetPad(2)->cd(2);
            hx = (TH1D*)fHResolution.ProjectionZ("Resolution");
            TPaveStats *stats = dynamic_cast<TPaveStats*>(hx->FindObject("stats"));
            if (stats)
            {
                stats->SetBit(BIT(17)); // TStyle.cxx: kTakeStyle=BIT(17)
                stats->SetX1NDC(0.63);
                stats->SetY1NDC(0.68);
            }

            hx->Fit("gaus", "Q", "", -1.0, 0.25);
            gPad=NULL;
            gStyle->SetOptFit(101);
        }
    }

    if (pad->GetPad(2))
    {
        pad->GetPad(2)->cd(1);
        UpdatePlot(fHEnergy, "yx", kTRUE);

        TLine *l = (TLine*)gPad->FindObject("TLine");
        if (l)
        {
            const Float_t min = TMath::Max(fHEnergy.GetXaxis()->GetXmin(), fHEnergy.GetYaxis()->GetXmin());
            const Float_t max = TMath::Min(fHEnergy.GetXaxis()->GetXmax(), fHEnergy.GetYaxis()->GetXmax());

            l->SetX1(min);
            l->SetX2(max);
            l->SetY1(min);
            l->SetY2(max);
        }

        pad->GetPad(2)->cd(2);
        UpdatePlot(fHResolution, "zy");

        pad->GetPad(2)->cd(3);
        UpdatePlot(fHResolution, "zx");
    }
}

void MHEnergyEst::UpdatePlot(TH3 &h, const char *how, Bool_t logy)
{
    TH2D *hyx=0;
    if (!(hyx=(TH2D*)gPad->FindObject(MString::Format("%s_%s", h.GetName(), how))))
        return;

    TH2D *h2 = (TH2D*)h.Project3D(MString::Format("dum_%s", how));
    hyx->Reset();
    hyx->Add(h2);
    delete h2;

    TH1D *hx = 0;
    if ((hx=(TH1D*)gPad->FindObject(MString::Format("Prof%s", h.GetName()))))
    {
        hx = hyx->ProfileX(Form("Prof%s", h.GetName()), -1, 9999, "s");

        if (logy && hx->GetMaximum()>0)
            gPad->SetLogy();
    }
}

TH1 *MHEnergyEst::MakePlot(TH3 &h, const char *how)
{
    gPad->SetBorderMode(0);
    gPad->SetLogx();
    gPad->SetGridx();
    gPad->SetGridy();

    // Results in crashes....
    //gROOT->GetListOfCleanups()->Add(gPad); // WHY?

    TH2D *h2 = (TH2D*)h.Project3D(how);
    h2->SetDirectory(NULL);
    h2->SetBit(kCanDelete);
    h2->SetFillColor(kBlue);
    h2->SetLineColor(kRed);

    TH1D *h1 = h2->ProfileX(Form("Prof%s", h.GetName()), -1, 9999, "s");
    h1->SetDirectory(NULL);
    h1->SetBit(kCanDelete);
    h1->SetLineWidth(2);
    h1->SetLineColor(kBlue);
    h1->SetFillStyle(4000);
    h1->SetStats(kFALSE);

    h2->Draw("");
    h1->Draw("E0 hist C same");
//    h1->Draw("Chistsame");

    return h2;
}


// --------------------------------------------------------------------------
//
// Draw the histogram
//
void MHEnergyEst::Draw(Option_t *opt)
{
    TVirtualPad *pad = gPad ? gPad : MakeDefCanvas(this);

    // Do the projection before painting the histograms into
    // the individual pads
    AppendPad("");

    pad->SetBorderMode(0);
    pad->Divide(2, 1, 1e-10, 1e-10);

    TH1 *h;

    pad->cd(1);
    gPad->SetBorderMode(0);

    gPad->Divide(1, 2, 1e-10, 1e-10);

    TVirtualPad *pad2 = gPad;

    pad2->cd(1);
    gPad->SetBorderMode(0);

    gPad->SetGridx();
    gPad->SetGridy();
    gPad->SetLogx();
    h = (TH1D*)fHEnergy.Project3D("ey");
    h->SetBit(TH1::kNoStats);
    h->SetTitle("Energy disribution: Monte Carlo E_{mc} (black), Estimated E_{est} (green)");
    h->SetXTitle("E [GeV]"); // E_mc
    h->SetYTitle("Counts");
    h->SetBit(kCanDelete);
    h->SetDirectory(NULL);
    h->SetMarkerStyle(kFullDotMedium);
    h->Draw();

    h = (TH1D*)fHEnergy.Project3D("ex");
    h->SetBit(TH1::kNoTitle|TH1::kNoStats);
    h->SetXTitle("E [GeV]"); // E_est
    h->SetYTitle("Counts");
    h->SetBit(kCanDelete);
    h->SetDirectory(NULL);
    h->SetMarkerStyle(kFullDotMedium);
    h->SetLineColor(kGreen);
    h->SetMarkerColor(kGreen);
    h->Draw("same");

    // FIXME: LEGEND

    pad2->cd(2);
    gPad->SetBorderMode(0);

    TVirtualPad *pad3 = gPad;
    pad3->Divide(2, 1, 1e-10, 1e-10);
    pad3->cd(1);
    gPad->SetBorderMode(0);
    gPad->SetGridx();
    gPad->SetGridy();
    h = fHEnergy.Project3D("ez");
    h->SetTitle("Zenith Angle Distribution");
    h->SetBit(TH1::kNoStats);
    h->SetDirectory(NULL);
    h->SetXTitle("\\Theta [\\circ]");
    h->SetBit(kCanDelete);
    h->Draw();

    pad3->cd(2);
    gPad->SetBorderMode(0);
    gPad->SetGridx();
    gPad->SetGridy();
    /*
    h = fHImpact.ProjectionX("Impact", -1, 9999, "e");
    h->SetBit(TH1::kNoStats);
    h->SetTitle("Distribution of Impact");
    h->SetDirectory(NULL);
    h->SetXTitle("Impact [m]");
    h->SetBit(kCanDelete);
    h->Draw();*/
    // ----------------------------------------
    h = fHResolution.ProjectionZ("Resolution");
    //h->SetXTitle("\\frac{lg(E_{est}) - lg(E_{mc})}{lg(E_{mc})}");
    h->SetYTitle("Counts");
    h->SetTitle("Distribution of \\Delta E/E");
    h->SetDirectory(NULL);
    h->SetBit(kCanDelete);
    h->GetXaxis()->SetRangeUser(-1.3, 1.3);
    h->Draw("");
    //h->Fit("gaus");
    // ----------------------------------------

    pad->cd(2);
    gPad->SetBorderMode(0);

    gPad->Divide(1, 3, 1e-10, 1e-10);
    pad2 = gPad;

    pad2->cd(1);
    h = MakePlot(fHEnergy, "yx");
    h->SetXTitle("E_{mc} [GeV]");
    h->SetYTitle("E_{est} [GeV]");
    h->SetTitle("Estimated energy E_{est} vs Monte Carlo energy E_{mc}");

    TLine line;
    line.DrawLine(0,0,1,1);

    pad2->cd(2);
    h = MakePlot(fHResolution, "zy");
    h->SetXTitle("E_{mc} [GeV]");
    h->SetYTitle("(E_{est}/E_{mc}-1");
    h->SetTitle("Energy resolution \\Delta E/E vs Monte Carlo Energy E_{mc}");
    h->SetMinimum(-1.3);
    h->SetMaximum(1.3);

    pad2->cd(3);
    h = MakePlot(fHResolution, "zx");
    h->SetXTitle("E_{est} [GeV]");
    h->SetYTitle("(E_{est}/E_{mc}-1");
    h->SetTitle("Energy resolution \\Delta E/E vs estimated Energy E_{est}");
    h->SetMinimum(-1.3);
    h->SetMaximum(1.3);
}

// --------------------------------------------------------------------------
//
// Returns the mapped value from the Matrix
//
Double_t MHEnergyEst::GetVal(Int_t i) const
{
    return (*fMatrix)[fMap[i]];
}

// --------------------------------------------------------------------------
//
// You can use this function if you want to use a MHMatrix instead of the
// given containers. This function adds all necessary columns to the
// given matrix. Afterward you should fill the matrix with the corresponding
// data (eg from a file by using MHMatrix::Fill). If you now loop
// through the matrix (eg using MMatrixLoop) MEnergyEstParam2::Process
// will take the values from the matrix instead of the containers.
//
void MHEnergyEst::InitMapping(MHMatrix *mat)
{
    if (fMatrix)
        return;

    fMatrix = mat;

    fMap[0] = fMatrix->AddColumn("MMcEvt.fEnergy");
    fMap[1] = fMatrix->AddColumn("MMcEvt.fImpact/100");
    fMap[2] = fMatrix->AddColumn("MMcEvt.fTelescopeTheta*kRad2Deg");
}

void MHEnergyEst::StopMapping()
{
    fMatrix = NULL; 
}