#include "TH1F.h"
#include "TH2F.h"
#include "TH3F.h"
#include "TLine.h"
#include "TTree.h"

#include "TCanvas.h"
#include "TStyle.h"
#include "TString.h"
#include "TFile.h"
#include "TMath.h"

#include "MReadTree.h"
#include "MSrcPosCam.h"
#include "MWriteRootFile.h"
#include "MParList.h"
#include "MTaskList.h"
#include "MEvtLoop.h"

#include "MHillas.h"
#include "MHillasExt.h"
#include "MHillasExt.h"
#include "MHillasSrc.h"
#include "MNewImagePar.h"

#include "MTaskInteractive.h"

#include "iostream.h"
#include "stdlib.h"
#include "stdio.h"
#include "string.h"

//--------------------------------------------------------------------------------------
// MAIN INPUT/OUTPUT

#include "IOMkn421.h"
//--------------------------------------------------------------------------------------

Double_t estposx =  0.;
Double_t estposy =  0.;

//--------------------------------------------------------------------------------------
// setting for hists
const Int_t nBinSize=200;
const Int_t nBinDist=1;
const Int_t nBinLW=1000;

// limits for size and dist for histograms
Double_t hSizeLo=2.;
Double_t hSizeUp=6.;
Double_t hDistLo=0.5;
Double_t hDistUp=1.15;

// limits for length, width for histograms
Double_t hLengthLo=0.;
Double_t hLengthUp=0.8;
Double_t hWidthLo =0.;
Double_t hWidthUp =0.6;

// limits for scaled Hillas for histograms
Double_t hSLenLo =0.;
Double_t hSLenUp =15.;
Double_t hSWidLo =0.;
Double_t hSWidUp =7.;

// special limits for hLength, hWidth (see ReadHillas())
const Int_t nBinLWFine=1000;
Double_t hLenFineLo=0.;
Double_t hLenFineUp=1.;
Double_t hWidFineLo=0.;
Double_t hWidFineUp=1.;

// hists for mean of unscaled Hillas
TH2F hLenMean("hLenMean","",nBinSize,hSizeLo,hSizeUp, nBinDist,hDistLo,hDistUp);
TH2F hWidMean("hWidMean","",nBinSize,hSizeLo,hSizeUp, nBinDist,hDistLo,hDistUp);

TH1F hLenMeanSize("hLenMeanSize","",nBinSize,hSizeLo,hSizeUp);
TH1F hLenMeanDist("hLenMeanDist","",nBinDist,hDistLo,hDistUp);
TH1F hWidMeanSize("hWidMeanSize","",nBinSize,hSizeLo,hSizeUp);
TH1F hWidMeanDist("hWidMeanDist","",nBinDist,hDistLo,hDistUp);

// hists for unscaled Hillas
TH2F hLenSize("hLenSize","",nBinSize,hSizeLo,hSizeUp, nBinLW,hLengthLo,hLengthUp);
TH2F hLenDist("hLenDist","",nBinDist,hDistLo,hDistUp, nBinLW,hLengthLo,hLengthUp);

TH2F hWidSize("hWidSize","",nBinSize,hSizeLo,hSizeUp, nBinLW,hWidthLo,hWidthUp);
TH2F hWidDist("hWidDist","",nBinDist,hDistLo,hDistUp, nBinLW,hWidthLo,hWidthUp);

// hists for mean-scaled Hillas
TH2F hSLenSize("hSLenSize","",nBinSize,hSizeLo,hSizeUp,nBinLW,hSLenLo,hSLenUp);
TH2F hSLenDist("hSLenDist","",nBinDist,hDistLo,hDistUp,nBinLW,hSLenLo,hSLenUp);

TH2F hSWidSize("hSWidSize","",nBinSize,hSizeLo,hSizeUp,nBinLW,hSWidLo,hSWidUp);
TH2F hSWidDist("hSWidDist","",nBinDist,hDistLo,hDistUp,nBinLW,hSWidLo,hSWidUp);


//****************************************************************************************
// first part: prepare for scaling, read Hillas to calc mean
//****************************************************************************************

void ReadHillas()
{
    // fill hillas into histograms to calc mean
    TFile *file=new TFile(fileHillas.Data());
    TTree *fEvents=(TTree*)file->Get("Events");

    // hists for Hillas
    // in following 2  hists the binning in length/width should be
    // very fine (like unbinned) and the limits should be generous
    TH3F hLength("hLength","",nBinSize,hSizeLo,hSizeUp, nBinDist,hDistLo,hDistUp, nBinLWFine,hLenFineLo,hLenFineUp);
    TH3F hWidth("hWidth",  "",nBinSize,hSizeLo,hSizeUp, nBinDist,hDistLo,hDistUp, nBinLWFine,hWidFineLo,hWidFineUp);

    TString strlen = Form("MHillas.fLength*%f:MHillasSrc.fDist*%f:log10(MHillas.fSize)",mm2deg,mm2deg);
    TString strwid = Form("MHillas.fWidth*%f :MHillasSrc.fDist*%f:log10(MHillas.fSize)",mm2deg,mm2deg);

    TString strcut="MHillas.fSize>0";
    fEvents->Project("hLength",strlen.Data(),strcut.Data());
    fEvents->Project("hWidth", strwid.Data(),strcut.Data());

    file->Close();

    // fill hists hLenMean, hWidMean (2-dim (size,dist)-hists)
    for(Int_t i=1;i<=nBinSize;i++)
        for(Int_t j=1;j<=nBinDist;j++)
        {
            Double_t lsum=0.; Double_t lsum2=0.; Double_t lcnt=0.;
            Double_t wsum=0.; Double_t wsum2=0.; Double_t wcnt=0.;

            for(Int_t k=1;k<=nBinLWFine;k++)
            {
                Double_t nl = hLength.GetBinContent(i,j,k);
                Double_t xl = hLength.GetZaxis()->GetBinCenter(k);
                
                lsum +=nl*xl;
                lsum2+=nl*xl*xl;
                lcnt +=nl;

                Double_t nw = hWidth.GetBinContent(i,j,k);
                Double_t xw = hWidth.GetZaxis()->GetBinCenter(k);

                wsum +=nw*xw;
                wsum2+=nw*xw*xw;
                wcnt +=nw;
            }

            Double_t lmean=0.;
            Double_t lmeanerr=0.;
            if(lcnt>1)lmean=lsum/lcnt;
            if(lcnt>2)lmeanerr=(lsum2/lcnt-lmean*lmean)/(lcnt-1.);
            lmeanerr=TMath::Sqrt(lmeanerr);

            hLenMean.SetBinContent(i,j,lmean);
            hLenMean.SetBinError(i,j,lmeanerr);

            Double_t wmean=0.;
            Double_t wmeanerr=0.;
            if(wcnt>1)wmean=wsum/wcnt;
            if(wcnt>2)wmeanerr=(wsum2/wcnt-wmean*wmean)/(wcnt-1.);
            wmeanerr=TMath::Sqrt(wmeanerr);

            hWidMean.SetBinContent(i,j,wmean);
            hWidMean.SetBinError(i,j,wmeanerr);
        }

    // fill hists hLenMeanSize, hWidMeanSize (1-dim hists, x-axis = size)
    for(Int_t i=1;i<=nBinSize;i++)
    {
        Double_t lsum=0.; Double_t lsum2=0.; Double_t lcnt=0.;
        Double_t wsum=0.; Double_t wsum2=0.; Double_t wcnt=0.;

        for(Int_t j=1;j<=nBinDist;j++)
            for(Int_t k=1;k<=nBinLWFine;k++)
            {
                Double_t nl = hLength.GetBinContent(i,j,k);
                Double_t xl = hLength.GetZaxis()->GetBinCenter(k);
                
                lsum +=nl*xl;
                lsum2+=nl*xl*xl;
                lcnt +=nl;

                Double_t nw = hWidth.GetBinContent(i,j,k);
                Double_t xw = hWidth.GetZaxis()->GetBinCenter(k);

                wsum +=nw*xw;
                wsum2+=nw*xw*xw;
                wcnt +=nw;
            }

        Double_t lmean=0.;
        Double_t lmeanerr=0.;
        if(lcnt>1)lmean=lsum/lcnt;
        if(lcnt>2)lmeanerr=(lsum2/lcnt-lmean*lmean)/(lcnt-1.);
        lmeanerr=TMath::Sqrt(lmeanerr);

        hLenMeanSize.SetBinContent(i,lmean);
        hLenMeanSize.SetBinError(i,lmeanerr);

        Double_t wmean=0.;
        Double_t wmeanerr=0.;
        if(wcnt>1)wmean=wsum/wcnt;
        if(wcnt>2)wmeanerr=(wsum2/wcnt-wmean*wmean)/(wcnt-1.);
        wmeanerr=TMath::Sqrt(wmeanerr);

        hWidMeanSize.SetBinContent(i,wmean);
        hWidMeanSize.SetBinError(i,wmeanerr);
    }

    // fill hists hLenMeanDist, hWidMeanDist (1-dim hists, x-axis = dist)
    for(Int_t j=1;j<=nBinDist;j++)
    {
        Double_t lsum=0.; Double_t lsum2=0.; Double_t lcnt=0.;
        Double_t wsum=0.; Double_t wsum2=0.; Double_t wcnt=0.;

        for(Int_t i=1;i<=nBinSize;i++)
            for(Int_t k=1;k<=nBinLWFine;k++)
            {
                Double_t nl = hLength.GetBinContent(i,j,k);
                Double_t xl = hLength.GetZaxis()->GetBinCenter(k);
                
                lsum +=nl*xl;
                lsum2+=nl*xl*xl;
                lcnt +=nl;

                Double_t nw = hWidth.GetBinContent(i,j,k);
                Double_t xw = hWidth.GetZaxis()->GetBinCenter(k);

                wsum +=nw*xw;
                wsum2+=nw*xw*xw;
                wcnt +=nw;
            }

        Double_t lmean=0.;
        Double_t lmeanerr=0.;
        if(lcnt>1)lmean=lsum/lcnt;
        if(lcnt>2)lmeanerr=(lsum2/lcnt-lmean*lmean)/(lcnt-1.);
        lmeanerr=TMath::Sqrt(lmeanerr);

        hLenMeanDist.SetBinContent(j,lmean);
        hLenMeanDist.SetBinError(j,lmeanerr);

        Double_t wmean=0.;
        Double_t wmeanerr=0.;
        if(wcnt>1)wmean=wsum/wcnt;
        if(wcnt>2)wmeanerr=(wsum2/wcnt-wmean*wmean)/(wcnt-1.);
        wmeanerr=TMath::Sqrt(wmeanerr);

        hWidMeanDist.SetBinContent(j,wmean);
        hWidMeanDist.SetBinError(j,wmeanerr);
    }

    return;
}


//****************************************************************************************
// second part: do scaling
//****************************************************************************************

// pointer to Hillas ParContainer
MHillas *fHillas=NULL;
MHillasSrc *fHillasSrc=NULL;


Int_t PreProcess(MParList *plist)
{
    fHillas = (MHillas*) plist->FindObject("MHillas");
    fHillasSrc = (MHillasSrc*) plist->FindObject("MHillasSrc");

    return (fHillas && fHillasSrc) ? kTRUE : kFALSE;
}

Int_t Process()
{
    // estimated source pos.
    MSrcPosCam srcpos;
    srcpos.SetXY(estposx*deg2mm,estposy*deg2mm);
    fHillasSrc->SetSrcPos(&srcpos);
    fHillasSrc->Calc(fHillas);

    Float_t width   = fHillas->GetWidth();
    Float_t length  = fHillas->GetLength();
    Float_t size    = fHillas->GetSize();
    Float_t meanx   = fHillas->GetMeanX();
    Float_t meany   = fHillas->GetMeanY();
    Float_t logsize = TMath::Log10(fHillas->GetSize());
    Float_t delta   = fHillas->GetDelta();
    Float_t dist    = fHillasSrc->GetDist();

    width *=mm2deg;
    length*=mm2deg;
    dist  *=mm2deg;

    hLenSize.Fill(logsize,length);
    hLenDist.Fill(dist,length);
    hWidSize.Fill(logsize,width);
    hWidDist.Fill(dist,width);

    // (log(size), dist) - bin
    Int_t ix=hLenMean.GetXaxis()->FindBin(logsize);
    Int_t iy=hLenMean.GetYaxis()->FindBin(dist);

    // check for over-/underflow
    /*if(ix*iy==0 || ix>nBinSize || iy>nBinDist)
    {
        cout<<"Warning, over-/underflow: Log(size)-bin="<<ix;
        cout<<" Dist-bin="<<iy<<endl<<flush;
    }*/
    if(ix==0)ix=1;if(ix==nBinSize+1)ix=nBinSize;
    if(iy==0)iy=1;if(iy==nBinDist+1)iy=nBinDist;

    Float_t lq=hLenMean.GetBinContent(ix,iy);
    Float_t wq=hWidMean.GetBinContent(ix,iy);

    if(lq>0.)lq=1./lq;
    if(wq>0.)wq=1./wq;

    width *=wq;
    length*=lq;

    hSLenSize.Fill(logsize,length);
    hSLenDist.Fill(dist,length);
    hSWidSize.Fill(logsize,width);
    hSWidDist.Fill(dist,width);

    // --------------------------------------------------------------------------
    //
    // This function is ment for special usage, please never try to set
    // values via this function
    //
    /* void MHillas::Set(const TArrayF &arr)
    {
        if (arr.GetSize() != 6)
            return;
    
        fLength = arr.At(0);  // [mm]        major axis of ellipse
        fWidth  = arr.At(1);  // [mm]        minor axis of ellipse
        fDelta  = arr.At(2);  // [rad]       angle of major axis with x-axis
        fSize   = arr.At(3);  // [#CerPhot]  sum of content of all pixels (number of Cherenkov photons)
        fMeanX  = arr.At(4);  // [mm]        x-coordinate of center of ellipse
        fMeanY  = arr.At(5);  // [mm]        y-coordinate of center of ellipse
    }*/

    TArrayF farr(6);
    farr[0]=length;
    farr[1]=width;
    farr[2]=delta;
    farr[3]=size;
    farr[4]=meanx;
    farr[5]=meany;

    fHillas->Set(farr);

    fHillas->SetReadyToSave();
    fHillasSrc->SetReadyToSave();

    return kTRUE;
}

Int_t PostProcess()
{

    return kTRUE;
}

void ScaleHillas(Int_t nent)
{
    //---------------------------------------------------------------
    // scale Hillas (length, width) to mean

    MParList plist;
    MTaskList tlist;
    plist.AddToList(&tlist);

    MReadTree read("Events", fileHillas);
    read.DisableAutoScheme();

    MTaskInteractive scale;
    scale.SetPreProcess(PreProcess);
    scale.SetProcess(Process);
    scale.SetPostProcess(PostProcess);

    MWriteRootFile write(fileHScaled.Data());
    write.AddContainer("MHillas",      "Events");
    write.AddContainer("MHillasExt",   "Events");
    write.AddContainer("MHillasSrc",   "Events");
    write.AddContainer("MNewImagePar", "Events");
    write.AddContainer("MRawRunHeader","Events");

    tlist.AddToList(&read);
    tlist.AddToList(&scale);
    tlist.AddToList(&write);

    MEvtLoop evtloop;
    evtloop.SetParList(&plist);

    if (!evtloop.Eventloop(nent))
        return;

    tlist.PrintStatistics();

    return;
}


//****************************************************************************************
// third part: check scaling
//****************************************************************************************

void CheckScaling()
{
    //---------------------------------------------------------------
    // check scaling
    TFile fileHist(fileHScalHist,"recreate");

    gStyle->SetPalette(1);


    TCanvas *c1=new TCanvas("c1","",800,400);
    if(nBinDist>1)
        c1->Divide(2,2);
    else
        c1->Divide(2,1);

    c1->cd(1);
    gPad->SetGrid();
    hSLenSize.SetStats(kFALSE);
    hSLenSize.GetXaxis()->SetTitle("log(size [#phot])");
    hSLenSize.GetYaxis()->SetTitle("scaled length");
    hSLenSize.SetLineWidth(2);
    hSLenSize.DrawCopy("colz");
    hSLenSize.Write();

    TLine lLen(2.7,1.,hSizeUp,1.);
    TLine lLenUp(2.7,LengthUp,hSizeUp,LengthUp);
    TLine lLenLo(2.7,LengthLo,hSizeUp,LengthLo);

    lLen.SetLineColor(2);
    lLen.SetLineWidth(2);
    lLenUp.SetLineWidth(2);
    lLenLo.SetLineWidth(2);
    lLen.DrawClone();
    lLenUp.DrawClone();
    lLenLo.DrawClone();


    c1->cd(2);
    gPad->SetGrid();
    hLenSize.SetStats(kFALSE);
    hLenSize.SetTitle("");
    hLenSize.GetXaxis()->SetTitle("log(size [#phot])");
    hLenSize.GetYaxis()->SetTitle("length [deg]");
    hLenSize.SetLineWidth(2);
    hLenSize.DrawCopy("colz");
    hLenSize.Write();
    hLenMeanSize.SetLineColor(2);
    hLenMeanSize.SetLineWidth(2);
    hLenMeanSize.DrawCopy("e same");
    hLenMeanSize.Write();

    TH1F hLenCutLo("hLenCutLo","",nBinSize,hSizeLo,hSizeUp);
    TH1F hLenCutUp("hLenCutUp","",nBinSize,hSizeLo,hSizeUp);
    for(Int_t i=1;i<=nBinSize;i++)
    {
        Float_t meanval=hLenMeanSize.GetBinContent(i);
        hLenCutLo.SetBinContent(i,meanval*LengthLo);
        hLenCutUp.SetBinContent(i,meanval*LengthUp);
    }
    hLenCutLo.SetLineWidth(2);
    hLenCutUp.SetLineWidth(2);
    hLenCutLo.DrawCopy("][ same");
    hLenCutUp.DrawCopy("][ same");


    if(nBinDist>1)
    {
        c1->cd(3);
        gPad->SetGrid();
        hSLenDist.SetStats(kFALSE);
        hSLenDist.GetXaxis()->SetTitle("dist [deg]");
        hSLenDist.GetYaxis()->SetTitle("scaled length");
        hSLenDist.SetLineWidth(2);
        hSLenDist.DrawCopy("colz");
        hSLenDist.Write();
    
        c1->cd(4);
        gPad->SetGrid();
        hLenDist.SetStats(kFALSE);
        hLenDist.SetTitle("");
        hLenDist.GetXaxis()->SetTitle("dist [deg]");
        hLenDist.GetYaxis()->SetTitle("length");
        hLenDist.SetLineWidth(2);
        hLenDist.DrawCopy("colz");
        hLenDist.Write();
        hLenMeanDist.SetLineColor(2);
        hLenMeanDist.SetLineWidth(2);
        hLenMeanDist.DrawCopy("e same");
        hLenMeanDist.Write();
    }

    TCanvas *c2=new TCanvas("c2","",800,400);
    if(nBinDist>1)
        c2->Divide(2,2);
    else
        c2->Divide(2,1);

    c2->cd(1);
    gPad->SetGrid();
    hSWidSize.SetStats(kFALSE);
    hSWidSize.GetXaxis()->SetTitle("log(size [#phot])");
    hSWidSize.GetYaxis()->SetTitle("scaled width");
    hSWidSize.SetLineWidth(2);
    hSWidSize.DrawCopy("colz");
    hSWidSize.Write();

    TLine lWid(2.7,1.,hSizeUp,1.);
    TLine lWidUp(2.7,WidthUp,hSizeUp,WidthUp);
    TLine lWidLo(2.7,WidthLo,hSizeUp,WidthLo);

    lWid.SetLineColor(2);
    lWid.SetLineWidth(2);
    lWidUp.SetLineWidth(2);
    lWidLo.SetLineWidth(2);
    lWid.DrawClone();
    lWidUp.DrawClone();
    lWidLo.DrawClone();

    c2->cd(2);
    gPad->SetGrid();
    hWidSize.SetStats(kFALSE);
    hWidSize.SetTitle("");
    hWidSize.GetXaxis()->SetTitle("log(size [#phot])");
    hWidSize.GetYaxis()->SetTitle("width [deg]");
    hWidSize.SetLineWidth(2);
    hWidSize.DrawCopy("colz");
    hWidSize.Write();
    hWidMeanSize.SetLineColor(2);
    hWidMeanSize.SetLineWidth(2);
    hWidMeanSize.DrawCopy("e same");
    hWidMeanSize.Write();

    TH1F hWidCutLo("hWidCutLo","",nBinSize,hSizeLo,hSizeUp);
    TH1F hWidCutUp("hWidCutUp","",nBinSize,hSizeLo,hSizeUp);
    for(Int_t i=1;i<=nBinSize;i++)
    {
        Float_t meanval=hWidMeanSize.GetBinContent(i);
        hWidCutLo.SetBinContent(i,meanval*WidthLo);
        hWidCutUp.SetBinContent(i,meanval*WidthUp);
    }
    hWidCutLo.SetLineWidth(2);
    hWidCutUp.SetLineWidth(2);
    hWidCutLo.DrawCopy("][ same");
    hWidCutUp.DrawCopy("][ same");

    if(nBinDist>1)
    {
        c2->cd(3);
        gPad->SetGrid();
        hSWidDist.SetStats(kFALSE);
        hSWidDist.GetXaxis()->SetTitle("dist [deg]");
        hSWidDist.GetYaxis()->SetTitle("scaled width");
        hSWidDist.SetLineWidth(2);
        hSWidDist.DrawCopy("colz");
        hSWidDist.Write();
    
        c2->cd(4);
        gPad->SetGrid();
        hWidDist.SetStats(kFALSE);
        hWidDist.SetTitle("");
        hWidDist.GetXaxis()->SetTitle("dist [deg]");
        hWidDist.GetYaxis()->SetTitle("width");
        hWidDist.SetLineWidth(2);
        hWidDist.DrawCopy("colz");
        hWidDist.Write();
        hWidMeanDist.SetLineColor(2);
        hWidMeanDist.SetLineWidth(2);
        hWidMeanDist.DrawCopy("e same");
        hWidMeanDist.Write();
    }

    // mean values
    if(nBinDist>1)
    {
        TCanvas *c4=new TCanvas("c4","",800,400);
        c4->Divide(1,2);
    
        c4->cd(1);
        gPad->SetGrid();
        hLenMean.SetStats(kFALSE);
        hLenMean.SetTitle("mean_length(log(size),dist)");
        hLenMean.GetXaxis()->SetTitle("log(size [#phot])");
        hLenMean.GetYaxis()->SetTitle("dist [deg])");
        hLenMean.DrawCopy("colz");
        hLenMean.Write();
    
        c4->cd(2);
        gPad->SetGrid();
        hWidMean.SetStats(kFALSE);
        hWidMean.SetTitle("mean_width(log(size),dist)");
        hWidMean.GetXaxis()->SetTitle("log(size [#phot])"); ;
        hWidMean.GetYaxis()->SetTitle("dist [deg])");
        hWidMean.DrawCopy("colz");
        hWidMean.Write("colz");
    }

    fileHist.Close();

    return;
}

//****************************************************************************************
// main
//****************************************************************************************

void RunScaling(Int_t nent=-1)
{
    printf("Prepare for scaling...");
    ReadHillas();
    printf("done.\n");

    ScaleHillas(nent);

    CheckScaling();

    return;
}