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

#include "MReadTree.h"

#include "MHillasSrc.h"
#include "MSrcPosCam.h"

#include "MHillasCalc.h"
#include "MHillasSrcCalc.h"
#include "MTaskInteractive.h"
#include "MContinue.h"
#include "MRawRunHeader.h"
#include "MWriteRootFile.h"

#include "TH2F.h"
#include "TF2.h"
#include "TArrayF.h"
#include "TObjArray.h"
#include "TStyle.h"
#include "TFile.h"
#include "TMath.h"
#include "TString.h"
#include "TCanvas.h"
#include "TGraphErrors.h"
#include "TMarker.h"

#include "iostream.h"

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

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

Double_t rmean=0.3;
Double_t dtlimit=10;  // time slice width [min]
//Double_t hlimit =0.85;
Double_t hlimit =0.9;

//----------------------------------------------------------------
// false source hist ---------------------------------------------
//----------------------------------------------------------------
const Int_t nx = 100; // no. of points in x-grid
const Int_t ny = 100; // no. of points in y-grid

Double_t alimit = 8.; // signal region in alpha-plot [degree]

Double_t xmin = -1.;  // units in degree !!
Double_t xmax =  1.;
Double_t ymin = -1.;
Double_t ymax =  1.;

Double_t dx=(xmax-xmin)/Double_t(nx);
Double_t dy=(ymax-ymin)/Double_t(ny);

TH2F FSrcHist("FSrcHist","",nx,xmin-dx,xmax+dx,ny,ymin-dy,ymax+dy);
TObjArray FSrcArray(1);

TH1F APlot("APlot","",36,0,90);
TObjArray APlotArray(1);

TArrayF xpeak(100);
TArrayF ypeak(100);

TArrayF xpeakerr(100);
TArrayF ypeakerr(100);

Int_t runno     = -1;
Int_t runno_old = -1;
Int_t runno_lo  = -1;
Int_t runno_up  = -1;

Double_t mjdend_old  = -1.;
Double_t mjdstart_lo = -1.;
Double_t mjdstart_up = -1.;

Double_t mjd_lo  = -1.;
Double_t mjd_up  = -1.;

Double_t dt  = 0.;
Int_t islice = 0;

TF2  *fpeak   = NULL;
FILE *fsrcpos = NULL;

//****************************************************************************************
// first part: track the source in time-intervalls of duration dtlimit
//****************************************************************************************

MHillas       *hil=NULL;
MRawRunHeader *run=NULL;

Double_t gaus2d(Double_t *x,Double_t *par)
{
    Double_t N=par[0];  //integral
    Double_t mx=par[1]; //mean_x
    Double_t my=par[2]; //mean_y
    Double_t sx=par[3]; //sigma_x
    Double_t sy=par[4]; //sigma_y
    Double_t rho=par[5];//correlation factor
    Double_t b=par[6];  //constant background

    Double_t z=(x[0]-mx)*(x[0]-mx)/sx/sx;
    z+=(x[1]-my)*(x[1]-my)/sy/sy;
    z-=2.*rho*(x[0]-mx)*(x[1]-my)/sx/sy;

    Double_t val=N/(2.*TMath::Pi()*sx*sy*sqrt(1-rho*rho));
    val*=exp( -z/(2.*(1.-rho*rho)) );
    val+=b;
    return val;
}

Bool_t CalcPeakXY(Double_t *xmean, Double_t *ymean,
                  Double_t *xmeanerr, Double_t *ymeanerr)
{
    Int_t mbin=FSrcHist.GetMaximumBin();
    Int_t biny = mbin/(nx+2);
    Int_t binx = mbin%(nx+2);
  
    if(mbin!=FSrcHist.GetBin(binx,biny))
    {
        cout<<"Error extracting binx, biny from global bin!"<<endl;
        cout<<" mbin="<<mbin<<" binfound="<<FSrcHist.GetBin(binx,biny)<<endl;
        cout<<" binx="<<binx<<" biny="<<biny<<endl;
  
        return kFALSE;
    }
  
    Double_t ndelta = rmean/FSrcHist.GetBinWidth(mbin);

    Double_t nmax=FSrcHist.GetBinContent(binx,biny);

    Int_t ilo=Int_t(binx-ndelta);
    Int_t iup=Int_t(binx+ndelta);

    Int_t jlo=Int_t(biny-ndelta);
    Int_t jup=Int_t(biny+ndelta);

    Double_t xsum  = 0.; Double_t xsum2 = 0.;
    Double_t ysum  = 0.; Double_t ysum2 = 0.;
    Double_t cnt   = 0.;
    //Double_t cov   = 0.;

    for(Int_t i=ilo;i<=iup;i++)
        for(Int_t j=jlo;j<=jup;j++)
        {
            Double_t n=(Double_t)FSrcHist.GetBinContent(i,j);
            if(n<hlimit*nmax)continue;
            Double_t x=(Double_t)FSrcHist.GetXaxis()->GetBinCenter(i);
            Double_t y=(Double_t)FSrcHist.GetYaxis()->GetBinCenter(j);

            cnt+=n;

            xsum +=n*x;
            xsum2+=n*x*x;
            ysum +=n*y;
            ysum2+=n*y*y;
        }

    Double_t xbar=0.; Double_t xbarerr=0.;
    Double_t ybar=0.; Double_t ybarerr=0.;

    if(cnt>1)
    {
        xbar=xsum/cnt;
        ybar=ysum/cnt;
    }
    if(cnt>2)
    {
        xbarerr=cnt*(xsum2/cnt-xbar*xbar)/(cnt-1.);
        xbarerr=TMath::Sqrt(xbarerr);

        ybarerr=cnt*(ysum2/cnt-ybar*ybar)/(cnt-1.);
        ybarerr=TMath::Sqrt(ybarerr);
    }

    /*
    //fit not yet tested!!
    Double_t b=FSrcHist.GetEntries()/nx/ny;
    Double_t nn=b*ndelta*ndelta;

    fpeak->SetParameter(0,nn);     // integral
    fpeak->SetParameter(1,*xmean); // mean_x
    fpeak->SetParameter(2,*ymean); // mean_y
    fpeak->SetParameter(3,0.15);   // sigma_x
    fpeak->SetParameter(4,0.15);   // sigma_y
    fpeak->SetParameter(5,0.);     // correlation factor
    fpeak->SetParameter(6,b);      // constant background

    FSrcHist.Fit("fpeak");
    xbar=fpeak->GetParameter(1);
    ybar=fpeak->GetParameter(2);
    xbarerr=fpeak->GetParameter(3);
    ybarerr=fpeak->GetParameter(4);
    */

    *xmean=xbar;
    *ymean=ybar;
    *xmeanerr=xbarerr;
    *ymeanerr=ybarerr;

    return kTRUE;
}

void AddFSrcHist(Int_t i,Double_t delta)
{
    Int_t n=FSrcArray.GetSize();
    if(i>=n) FSrcArray.Expand(n+1);

    TH2F *htemp=(TH2F*)FSrcHist.Clone();
    TString name=Form("%d",i);
    TString title=Form("TimeIntervall = %f min",delta);

    htemp->SetName(name.Data());
    htemp->SetTitle(title.Data());
    htemp->SetDrawOption("colz");

    FSrcArray[i] = htemp;

    return;
}

void AddAPlot(Int_t i,Double_t delta)
{
    Int_t n=APlotArray.GetSize();
    if(i>=n) APlotArray.Expand(n+1);

    TH1F *htemp=(TH1F*)APlot.Clone();
    TString name=Form("%d",i);
    TString title=Form("TimeIntervall = %f min",delta);

    htemp->SetName(name.Data());
    htemp->SetTitle(title.Data());
    htemp->SetFillColor(16);

    APlotArray[i] = htemp;

    return;
}

//----------------------------------------------------------------
// Source Position Grid ------------------------------------------
//----------------------------------------------------------------

Int_t STrackPreProc(MParList *plist)
{
    hil = (MHillas*) plist->FindObject("MHillas");
    run = (MRawRunHeader*) plist->FindObject("MRawRunHeader");

    return (hil && run) ? kTRUE : kFALSE;
}

Int_t STrackProc()
{
    runno  = run->GetRunNumber();
    if(runno==0) return kTRUE;

    Double_t mjdstart = (run->GetRunStart()).GetMjd();
    Double_t mjdend   = (run->GetRunEnd()).GetMjd();
    Double_t t = (mjdend-mjdstart)*24.*60.;

    Double_t dtrun;
    if(runno!=runno_old)
    {
        if(dt<0.001) {runno_lo=runno;mjd_lo=mjdstart;mjdend_old=mjdstart;} // very first run

        dtrun=(mjdend_old-mjd_lo)*24.*60.;
        //if(dt>dtlimit)
        if(dtrun>dtlimit)
        {
            runno_up=runno_old;
            mjd_up=mjdend_old;

            Double_t x; Double_t xerr;
            Double_t y; Double_t yerr;

            if(!CalcPeakXY(&x,&y,&xerr,&yerr)) return kFALSE;
            AddFSrcHist(islice,dt);

            xpeak[islice]=x;
            ypeak[islice]=y;
            xpeakerr[islice]=xerr;
            ypeakerr[islice]=yerr;

            printf("\n TimeSlice %d   RunNo %d-%d  Mjd %f-%f  Dt=%f \n",islice,runno_lo,runno_up,mjd_lo,mjd_up,dt);
            printf(" X=%f+-%f  Y=%f+-%f \n\n",x,xerr,y,yerr);
            fprintf(fsrcpos,"%d %d %d %f %f %f %f %f %f %f\n",islice,runno_lo,runno_up,x,y,xerr,yerr,mjd_lo,mjd_up,dt);

            FSrcHist.Reset();
            dt=0.;

            runno_lo=runno;
            mjd_lo=mjdstart;

            islice++;
        }

        dt+=t;
        printf("Runno=%d  MjdStart=%f  MjdEnde=%f dt=%f \n",runno,mjdstart,mjdend,dt);
        runno_old=runno;
        mjdend_old=mjdend;
    }


    MHillasSrc hillas_src;
    MSrcPosCam srcpos;

    Double_t length=hil->GetLength();
    Double_t width=hil->GetWidth();

    // size cut
    if(hil->GetSize()<SizeLo) return kTRUE;

    // cuts in scaled length, width
    if((length>LengthUp) || (length<LengthLo) || (width>WidthUp) || (width<WidthLo))
        return kTRUE;

    // cut in dist from camera center
    srcpos.SetXY(0,0);
    hillas_src.SetSrcPos(&srcpos);
    hillas_src.Calc(hil);

    Double_t dist=hillas_src.GetDist()*mm2deg;
    if(dist>DistUp || dist<DistLo) return kTRUE;

    // false source grid
    for (Int_t i=0;i<nx;i++)
        for (Int_t j=0;j<ny;j++)
        {
            Double_t x=(xmax-xmin)*Double_t(i)/Double_t(nx-1)+xmin;
            Double_t y=(ymax-ymin)*Double_t(j)/Double_t(ny-1)+ymin;

            srcpos.SetXY(x*deg2mm,y*deg2mm);
            hillas_src.SetSrcPos(&srcpos);
            hillas_src.Calc(hil);

            if(TMath::Abs(hillas_src.GetAlpha())<alimit)
                FSrcHist.Fill(x,y);

        }

    return kTRUE;
}

Int_t STrackPostProc()
{
    if(dt>0.001)
    {
        runno_up=runno_old;
        mjd_up=mjdend_old;

        Double_t x; Double_t xerr;
        Double_t y; Double_t yerr;

        if(!CalcPeakXY(&x,&y,&xerr,&yerr)) return kFALSE;
        AddFSrcHist(islice,dt);

        xpeak[islice]=x;
        ypeak[islice]=y;
        xpeakerr[islice]=xerr;
        ypeakerr[islice]=yerr;

        printf("\n TimeSlice %d   RunNo %d-%d  Mjd %f-%f  Dt=%f \n",islice,runno_lo,runno_up,mjd_lo,mjd_up,dt);
        printf(" X=%f+-%f  Y=%f+-%f \n\n",x,xerr,y,yerr);
        fprintf(fsrcpos,"%d %d %d %f %f %f %f %f %f %f\n",islice,runno_lo,runno_up,x,y,xerr,yerr,mjd_lo,mjd_up,dt);

        islice++;
    }

    return kTRUE;
}

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

void SourceTrack(Int_t nent)
{

    MParList plist;

    MTaskList tlist;
    plist.AddToList(&tlist);

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

    //TString runcuts("(MRawRunHeader.fRunNumber<17427)");
    //runcuts+="|| (MRawRunHeader.fRunNumber>17428)";
    //MContinue runsel(runcuts.Data());

    MTaskInteractive strack;
    strack.SetPreProcess(STrackPreProc);
    strack.SetProcess(STrackProc);
    strack.SetPostProcess(STrackPostProc);

    tlist.AddToList(&read);
    //tlist.AddToList(&runsel);
    tlist.AddToList(&strack);

    MEvtLoop evtloop;
    evtloop.SetParList(&plist);

    //------------------------------------------
    // initializations

    FSrcHist.SetDrawOption("colz");
    FSrcHist.SetStats(kFALSE);

    fpeak=new TF2("fpeak",gaus2d,-1,1,-1,1,6);
    fsrcpos=fopen(fileSrcPos.Data(),"w");

    //------------------------------------------
    // standard eventloop

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

    tlist.PrintStatistics();

    fclose(fsrcpos);

    TFile file(fileFSrcPlot.Data(),"recreate");
    FSrcArray.Write();
    FSrcHist.Write();
    file.Close();

    const Int_t dim=islice;
    TGraphErrors g(dim,xpeak.GetArray(),ypeak.GetArray(),
                   xpeakerr.GetArray(),ypeakerr.GetArray());

    //------------------------------------------
    // canvas with complete hist
    TCanvas *c1=new TCanvas("c1","",400,800);
    gStyle->SetPalette(1);

    c1->Divide(2,1);
    c1->cd(1);
    gPad->Update();

    FSrcHist.Reset();
    for(Int_t i=0;i<dim;i++)
    {
        TH2F *hptr = (TH2F*) (FSrcArray[i]);
        FSrcHist.Add(hptr);
    }

    TString label=Form("Size>%f  ",SizeLo);
    label+=Form("%f<Dist<%f  ",  DistLo,DistUp);
    label+=Form("%f<Width<%f  " ,WidthLo,WidthUp);
    label+=Form("%f<Length<%f  ",LengthLo,LengthUp);
    label+=Form("|Alpha|<%f"  ,alimit);

    FSrcHist.SetTitle(label.Data());
    FSrcHist.GetXaxis()->SetTitle("X");
    FSrcHist.GetYaxis()->SetTitle("Y");

    FSrcHist.DrawClone("colz");
    g.DrawClone("pl");

    c1->cd(2);
    gPad->Update();
    g.GetXaxis()->SetRangeUser(-1.,1.);
    g.GetYaxis()->SetRangeUser(-1.,1.);
    g.DrawClone("apl");

    //------------------------------------------
    // canvas with slides
    TCanvas *c2=new TCanvas("c2","",800,800);
    gStyle->SetPalette(1);

    Int_t n=Int_t(TMath::Ceil(sqrt(double(dim))));

    c2->Divide(n,n);
    for(Int_t i=0;i<dim;i++)
    {
        Int_t col=i/n;
        Int_t row=i%n;
        c2->cd(row*n+col+1);

        gPad->Update();
        TH2F *hptr = (TH2F*) (FSrcArray[i]);

        TMarker m;
        m.SetMarkerStyle(2);
        m.SetMarkerSize(4);
        m.SetMarkerColor(1);
        m.SetX(xpeak[i]);
        m.SetY(ypeak[i]);

        hptr->DrawClone("colz");
        m.DrawClone();
    }

    return;
}

//****************************************************************************************
// second part: correct for moving source position
//****************************************************************************************

// initilizations

Double_t dtref=0.;
Int_t islice_old;
Double_t dtref_old=0;

Double_t xref = 0.; Double_t xreferr = 0.;
Double_t yref = 0.; Double_t yreferr = 0.;

MSrcPosCam *src=NULL;

//----------------------------------------------------------------
// Source Position Grid ------------------------------------------
//----------------------------------------------------------------
Int_t SCorrectPreProc(MParList *plist)
{
    hil = (MHillas*) plist->FindObject("MHillas");
    src = (MSrcPosCam*) plist->FindCreateObj("MSrcPosCam");
    run = (MRawRunHeader*) plist->FindObject("MRawRunHeader");

    return (hil && run && src) ? kTRUE : kFALSE;
}

Int_t SCorrectProc()
{
    if(hil->GetSize()<=0.) return kTRUE;

    Int_t runno    = run->GetRunNumber();
    if(runno==0) return kTRUE;

    Double_t mjdstart = (run->GetRunStart()).GetMjd();
    Double_t mjdend   = (run->GetRunEnd()).GetMjd();
    Double_t t=(mjdend-mjdstart)*24.*60.;

    if(runno!=runno_old)
    {
        if((runno<runno_lo)||(runno>runno_up))
        {
            dt=t;

            fscanf(fsrcpos,"%d %d %d %lf %lf %lf %lf %lf %lf %lf\n",&islice,&runno_lo,&runno_up,&xref,&yref,&xreferr,&yreferr,&mjd_lo,&mjd_up,&dtref);
            printf("\nApplying correction for runs %d - %d (MJD %f - %f), DtRef=%f :\n",runno_lo,runno_up,mjd_lo,mjd_up,dtref);
            printf("Source pos set to X=%f  Y=%f \n\n",xref,yref);

            if(islice>0)
            {
                AddAPlot(islice_old,dtref_old);
                APlot.Reset();
            }
            islice_old=islice;
            dtref_old=dtref;
        }
        else
        {
            dt+=t;
        }
        printf("Runno=%d  MjdStart=%f  MjdEnde=%f Dt=%f \n",runno,mjdstart,mjdend,dt);
        runno_old=runno;

    }

    src->SetXY(xref*deg2mm,yref*deg2mm);
    src->SetReadyToSave();
    // false source grid

    MHillasSrc hillas_src;
    MSrcPosCam srcpos;

    Double_t length=hil->GetLength();
    Double_t width =hil->GetWidth();

    // size cut
    if(hil->GetSize()<SizeLo) return kTRUE;

    // cuts in scaled length, width
    if((length>LengthUp) || (length<LengthLo) || (width>WidthUp) || (width<WidthLo))
        return kTRUE;

    // cut in dist from camera center
    srcpos.SetXY(0,0);
    hillas_src.SetSrcPos(&srcpos);
    hillas_src.Calc(hil);

    Double_t dist=hillas_src.GetDist()*mm2deg;
    if(dist>DistUp || dist<DistLo) return kTRUE;

    // fill 1-dim APlot for extracted source position
    srcpos.SetXY(xref*deg2mm,yref*deg2mm);
    hillas_src.SetSrcPos(&srcpos);
    hillas_src.Calc(hil);
    APlot.Fill(TMath::Abs(hillas_src.GetAlpha()));

    // false source plot for corrected source movement
    for (Int_t i=0;i<nx;i++)
        for (Int_t j=0;j<ny;j++)
        {
            Double_t x=(xmax-xmin)*Double_t(i)/Double_t(nx-1)+xmin;
            Double_t y=(ymax-ymin)*Double_t(j)/Double_t(ny-1)+ymin;

            //printf("xref=%f  yref=%f \n",xref,yref);
            srcpos.SetXY((x+xref)*deg2mm,(y+yref)*deg2mm);
            hillas_src.SetSrcPos(&srcpos);
            hillas_src.Calc(hil);

            if(TMath::Abs(hillas_src.GetAlpha())<alimit)
                FSrcHist.Fill(x,y);
        }

    return kTRUE;
}

Int_t SCorrectPostProc()
{
    AddAPlot(islice_old,dtref_old);

    return kTRUE;
}

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

void SourceCorrect(Int_t nent)
{
    MParList plist;
    MTaskList tlist;
    plist.AddToList(&tlist);

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

    MTaskInteractive scorrect;
    scorrect.SetPreProcess(SCorrectPreProc);
    scorrect.SetProcess(SCorrectProc);
    scorrect.SetPostProcess(SCorrectPostProc);

    MHillasSrcCalc scalc;

    MWriteRootFile write(fileHScaledC.Data());
    write.AddContainer("MSrcPosCam",   "Events");
    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(&scorrect);
    tlist.AddToList(&scalc);
    tlist.AddToList(&write);

    MEvtLoop evtloop;
    evtloop.SetParList(&plist);

    //------------------------------------------
    // standard eventloop

    fsrcpos=fopen(fileSrcPos.Data(),"r");
    if(!fsrcpos)
    {
        cout<<"fsrcpos does not exist!"<<endl;
        return;
    }

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

    tlist.PrintStatistics();

    fclose(fsrcpos);


    gStyle->SetPalette(1);
    TFile file(fileFSrcPlotC.Data(),"recreate");
    FSrcHist.Write();
    APlotArray.Write();
    file.Close();

    //------------------------------------------
    // canvas with complete hist
    TCanvas *c=new TCanvas("c","",800,400);
    c->cd();

    TString label=Form("Size>%f  ",SizeLo);
    label+=Form("%f<Dist<%f  ",  DistLo,DistUp);
    label+=Form("%f<Width<%f  " ,WidthLo,WidthUp);
    label+=Form("%f<Length<%f  ",LengthLo,LengthUp);
    label+=Form("|Alpha|<%f"  ,alimit);

    FSrcHist.SetTitle(label.Data());
    FSrcHist.GetXaxis()->SetTitle("X");
    FSrcHist.GetYaxis()->SetTitle("Y");

    FSrcHist.DrawCopy("colz");

    //------------------------------------------
    // canvas with alpha hists
    TCanvas *c2=new TCanvas("c2","",800,800);
    gStyle->SetPalette(1);

    const Int_t dim=islice+1;
    Int_t n=Int_t(TMath::Ceil(sqrt(double(dim))));

    c2->Divide(n,n);
    for(Int_t i=0;i<dim;i++)
    {
        Int_t col=i/n;
        Int_t row=i%n;
        c2->cd(row*n+col+1);

        gPad->Update();
        TH1F *hptr = (TH1F*) (APlotArray[i]);

        hptr->DrawClone();
    }

    return;
}

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

void SrcCorrect(Int_t icorrect,Int_t nent=-1)
{

    if(!icorrect)
    {
        // (re-) initilizations
        islice=0;
        runno_old=-1;
        runno_lo=-1;
        runno_up=-1;

        dt=0;

        mjd_lo=-1;
        mjd_up=-1;

        hil=NULL; run=NULL;

        cout<<"Starting source track reconstruction."<<endl<<flush;
        FSrcHist.Reset();
        SourceTrack(nent);

    }else{

        // (re-) initilizations
        islice=0;
        runno_old=-1;
        runno_lo=-1;
        runno_up=-1;

        dt=0;

        mjd_lo=-1;
        mjd_up=-1;

        hil=NULL; run=NULL;

        cout<<"Starting source correct."<<endl<<flush;
        FSrcHist.Reset();
        SourceCorrect(nent);
    }

    return;
}
