/* ======================================================================== *\
!
! *
! * 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 <TLine.h>
#include <TCanvas.h> 
#include <TProfile.h>

#include "MString.h"

#include "MMcEvt.hxx"

#include "MEnergyEst.h"
#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 : "3-D histogram   E-true E-est Theta";

    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 \\frac{\\Delta E}{E_{mc}} vs E_{est} and E_{mc}");
    fHResolution.SetXTitle("E_{est} [GeV]");
    fHResolution.SetYTitle("E_{mc} [GeV]");
    fHResolution.SetZTitle("\\frac{E_{est} - E_{mc}}{E_{mc}}");

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

    fHEnergy.Sumw2();
    fHImpact.Sumw2();
    fHResolution.Sumw2();

    MBinning binsi, binse, binst, binsr;
    binse.SetEdgesLog(25, 10, 10000);
    binst.SetEdgesCos(50, 0, 60);
    binsi.SetEdges(10, 0, 400);
    binsr.SetEdges(10, 0, 1);

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

// --------------------------------------------------------------------------
//
// 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 = (MEnergyEst*)plist->FindObject("MEnergyEst");
    if (!fEnergy)
    {
        *fLog << err << "MEnergyEst 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;

    return kTRUE;
}

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

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

    fChisq += res*res;

    return kTRUE;
}

Bool_t MHEnergyEst::Finalize()
{
    fChisq /= GetNumExecutions();

    fResult->SetVal(fChisq);

    *fLog << all << "Mean Energy Resoltuion: " << Form("%.1f%%", TMath::Sqrt(fChisq)*100) << endl;

    return kTRUE;
}

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);
            /*h =*/ fHImpact.ProjectionX("Impact", -1, 9999, "e");

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

    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::Form("%s_%s", h.GetName(), how))))
        return;

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

    TH1D *hx = 0;
    if ((hx=(TH1D*)gPad->FindObject("Prof")))
    {
        hx = hyx->ProfileX("Prof", -1, 9999, "s");

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

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

    gROOT->GetListOfCleanups()->Add(gPad); // WHY?

    TH2D *h2 = (TH2D*)h.Project3D(how);
    TH1D *h1 = h2->ProfileX("Prof", -1, 9999, "s");

    h1->SetDirectory(NULL);
    //h1->SetBit(kCanDelete);
    h1->SetLineWidth(2);
    h1->SetLineColor(kRed); // PROBLEM!
    h1->SetStats(kFALSE);

    h2->SetDirectory(NULL);
    h2->SetBit(kCanDelete);
    h2->SetFillColor(kBlue);

    h1->Draw("E3");
    h2->Draw("boxsame");
    h1->Draw("Chistsame");

    return h1;
}


// --------------------------------------------------------------------------
//
// 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, 0, 0);

    TH1 *h;

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

    gPad->Divide(1, 2, 0, 0);

    TVirtualPad *pad2 = gPad;

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

    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, 0, 0);
    pad3->cd(1);
    gPad->SetBorderMode(0);
    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();

    pad3->cd(2);
    gPad->SetBorderMode(0);
    h = fHEnergy.Project3D("ez");
    h->SetTitle("Distribution of Theta");
    h->SetBit(TH1::kNoStats);
    h->SetDirectory(NULL);
    h->SetXTitle("\\Theta [\\circ]");
    h->SetBit(kCanDelete);
    h->Draw();

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

    gPad->Divide(1, 3, 0, 0);
    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("\\frac{E_{est} - E_{mc}}{E_{mc}}");
    h->SetTitle("Energy resolution \\Delta E / E vs Monte Carlo energy E_{mc}");
    h->SetMinimum(0);
    h->SetMaximum(1);

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

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