//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//   Plot parameters of reconstructed tracks
// This macro plots number of the hits contained in a track, the number of tracks per event, the (biased) residuals in y and z, its chi2/ndf and p value distributions, the position and angle of the track's staring point and a beam profile
// load macro with .L <macroname>.C+ to compile it and then enter ExecuteTrackAnalysis and press <tab> to see the arguments
// runnumber: which run? datanumber: which reconstruction of this run? nevents: process how many events? singletrackevents: only process single track events? dirname: output directory
//
// Environment:
//      GEM ALICE IROC prototype data analysis in fopiroot
//
// Author List:
//      Philipp Gadow          (philipp.gadow@mytum.de)
//
//
//-----------------------------------------------------------

#include <TClonesArray.h>
#include "TCanvas.h"
#include "TAxis.h"
#include "TFrame.h"
#include "TPad.h"
#include "TLatex.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TLegend.h"
#include "TF1.h"
#include "TH1F.h"
#include "TH2D.h"
#include "TROOT.h"
#include "TVector3.h"

#include <string>
#include <vector>
#include <map>
#include <algorithm>
#include <iostream>
#include <sstream>


#include <TpcEvent.h>
#include <TrackFitStat.h>
#include <GFTrack.h>

//---------------------------------------------
//Helper functions
//---------------------------------------------

string inttostring(Int_t input)
{
    string s;
    stringstream out;
    out << input;
    s = out.str();
    return s;
}


// Definition of Fit Functions (used for residual fit)
//----------------------------------------------------------
// Function for the very sharp peak / the asymmety of the peak
Double_t Gauss_0(Double_t *x, Double_t *par) {

    Double_t arg = 0;
    if (par[2]) arg = (x[0] - par[1])/par[2];

    Double_t sig = par[0]*TMath::Exp(-0.5*arg*arg);
    return sig;
}


// Function for the narrow Gaussian
Double_t Gauss_1(Double_t *x, Double_t *par) {

    Double_t arg = 0;
    if (par[2]) arg = (x[0] - par[1])/par[2];

    Double_t sig = par[0]*TMath::Exp(-0.5*arg*arg);
    return sig;
}

// Function for the broad Gaussian
Double_t Gauss_2(Double_t *x, Double_t *par) {

    Double_t arg = 0;
    if (par[2]) arg = (x[0] - par[1])/par[2];

    Double_t sig = par[0]*TMath::Exp(-0.5*arg*arg);
    return sig;
}


// Sum of three Gaussians
Double_t fitFunction(Double_t *x, Double_t *par) {
  return Gauss_0(x,par) + Gauss_1(x,&par[3]) + Gauss_2(x,&par[6]);
}
//----------------------------------------------------------


void ExecuteTrackAnalysis(bool verbose, int runnumber, int datanumber, int nevents, bool singletrackevents, bool pretty = false, TString dirname = "../plots/trkanalysis/"){

    if (pretty){
        gROOT->ProcessLine(".x ~/rootlogon_Bernhard.C");
        gROOT->SetStyle("col");
        gROOT->ForceStyle();
    }

    TString indirname = "/nfs/mds/data/tpc/alice/ps_test_beam/RecoOut/";
    TString infilename = "merge_" + inttostring(runnumber) + "_" + inttostring(datanumber) + "_corrected_hough_smoothed.reco.root";
    TString inpathname = indirname + infilename;

    // open file
    TFile *reco = new TFile(inpathname);
    // create canvas
    TCanvas *c1 = new TCanvas("c1","hits per track",1000,700);
    c1->SetFillColor(00);
    c1->SetGrid();

    TCanvas *c2 = new TCanvas("c2","tracks per event",1000,700);
    c2->SetFillColor(00);
    c2->SetGrid();

    TCanvas *c3 = new TCanvas("c3","residuals",1000,700);
    c3->SetFillColor(00);
    c3->Divide(2,2);
    c3->SetGrid();

    TCanvas *c4 = new TCanvas("c4","track statistics",1000,700);
    c4->SetFillColor(00);
    c4->Divide(2,1);
    c4->SetGrid();

    TCanvas *c5 = new TCanvas("c5","track parameters",1000,700);
    c5->SetFillColor(00);
    c5->Divide(3,2);
    c5->SetGrid();

    TCanvas *c6 = new TCanvas("c6","beam profile",1000,700);
    c6->SetFillColor(00);
    c6->SetGrid();

    // create histograms
    TH1F *histClustersTrack = new TH1F("histClustersTrack", "clusters per track", 100,0,100);
    TH1F *histTracksperEvent = new TH1F("histTracksperEvent", "tracks per event", 10,0,10);

    TH1F *histResidualsY = new TH1F("histhistResidualsY", "residuals y", 10000,-5,5);
    TH1F *histResidualsZ = new TH1F("histhistResidualsZ", "residuals Z", 10000,-5,5);

    TH1F *histChi2perNDF = new TH1F("histChi2perNDF", "Chi2/NDF", 3000,0,30);
    TH1F *histPVal = new TH1F("histPVal", "p value", 1000,0,1.2);

    TH1F *histtrkTheta = new TH1F("histtrkTheta", "Theta plane", 3000,-1.5,1.5);
    TH1F *histtrkPhi = new TH1F("histtrkPhi", "Phi angle", 3000,-1.5,1.5);
    TH1F *histtrkMom = new TH1F("histtrkMom", "momentum distribution of tracks", 700,0,7);
    TH1F *histtrkPosx = new TH1F("histtrkPosx", "x-position of track", 2000,0,200);
    TH1F *histtrkPosy = new TH1F("histtrkPosy", "y-position of track", 2000,-5,15);
    TH1F *histtrkPosz = new TH1F("histtrkPosz", "z-position of track", 1200,0,12);

    TH2D *histBeamProfile = new TH2D("histBeamProfile", "Beam Profile", 200,-5,15,120,0,12);

    // create Tree
    TTree *tpcTree = (TTree*)reco->Get("cbmsim");

    //save time and load only necessary branches
    tpcTree->SetBranchStatus("*",0);
    tpcTree->SetBranchStatus("TrackPostFit.*",1);
    tpcTree->SetBranchStatus("TrackFitStat_0.*",1);

    // create arrays
    TClonesArray *tracks = new TClonesArray("GFTrack");
    TClonesArray *fitstat = new TClonesArray("TrackFitStat");


    // set branches
    tpcTree->SetBranchAddress("TrackFitStat_0",&fitstat);
    tpcTree->SetBranchAddress("TrackPostFit",&tracks);


    // loop over events

    if (nevents == 0){
        nevents = tpcTree->GetEntries();
    }
    for( Int_t ev = 0; ev < nevents; ev++){

        if(ev%1000==0){
            cout<<"#################### Event done: "<<ev<<"##############"<<endl;
        }
        tracks->Delete();
        fitstat->Delete();
        tpcTree->GetEvent(ev);

        int numberoftracks = tracks->GetEntries();

        if (singletrackevents){
            if(numberoftracks!=1)
                continue;
        }


        // track analysis
        histTracksperEvent->Fill(numberoftracks);
        if(verbose){
            std::cout << "number of tracks per event:" << numberoftracks << std::endl;
        }

        //loop over tracks
        for(Int_t itr=0;itr<numberoftracks;++itr){
            GFTrack* trk=(GFTrack*)tracks->At(itr);
            GFTrackCand cand = trk->getCand();

            GFAbsTrackRep* rep = trk->getCardinalRep();

            if(rep->getStatusFlag()!=0){
                continue;
            }


            unsigned int numberofhits = cand.getNHits();

            double pVal = trk->getCardinalRep()->getPVal();

            if(verbose)
                std::cout << "number of hits was obtained "<< numberofhits<< std::endl;


            TVector3 normalXZplane (0,1,0);
            TVector3 normalXYplane (0,0,1);

            double tracktheta = TMath::ASin(trk->getCardinalRep()->getMom().Dot(normalXYplane)/trk->getCardinalRep()->getMom().Mag());
            double trackphi = TMath::ASin(trk->getCardinalRep()->getMom().Dot(normalXZplane)/trk->getCardinalRep()->getMom().Mag());

            double trackposx = trk->getCardinalRep()->getPos().X();
            double trackposy = trk->getCardinalRep()->getPos().Y();
            double trackposz = trk->getCardinalRep()->getPos().Z();
            double trackmom = trk->getCardinalRep()->getMom().Mag();

            histPVal->Fill(pVal);
            histClustersTrack->Fill(numberofhits);
            histtrkTheta->Fill(tracktheta);
            histtrkPhi->Fill(trackphi);
            histtrkMom->Fill(trackmom);
            histtrkPosx->Fill(trackposx);
            histtrkPosy->Fill(trackposy);
            histtrkPosz->Fill(trackposz);
            histBeamProfile->Fill(trackposy,trackposz);
        }
        //residuals
        if (!(fitstat->GetEntriesFast()<1)){

            TrackFitStat* stat = (TrackFitStat*) fitstat->At(0);
            const std::vector<double>* resY = stat->GetResY();
            const std::vector<double>* resZ = stat->GetResZ();

            for(unsigned int r=0; r<resY->size(); r++){
                histResidualsY->Fill(resY->at(r));
            }
            for(unsigned int r=0; r<resZ->size(); r++){
                histResidualsZ->Fill(resZ->at(r));
            }

            double chi2 = stat->getChi2();
            int NDF = stat->getNDF();

            histChi2perNDF->Fill(chi2/float(NDF));

        }
    }


    //define output filenames
    TString filename = "merge_" + inttostring(runnumber) + "_" + inttostring(datanumber);
    if(singletrackevents){
        filename += "_singletrackevents_";
    }
    TString savefilename = dirname + filename + "trackanalysiszpospeak.root";

    TFile *savefile = new TFile(savefilename.Data(),"RECREATE");

    //  Draw the histograms
    c1->cd();
    histClustersTrack->GetXaxis()->SetTitle("clusters per track");
    histClustersTrack->GetYaxis()->SetTitle("# of entries");
    histClustersTrack->Draw();
    histClustersTrack->Write();
    c1->Write();

    c2->cd();
    histTracksperEvent->GetXaxis()->SetTitle("tracks per event");
    histTracksperEvent->GetXaxis()->CenterLabels();
    histTracksperEvent->GetYaxis()->SetTitle("# of entries");
    histTracksperEvent->Draw();
    histTracksperEvent->Write();
    c2->Write();


    //--------------------------------------------------------------------------------------------------------------------------------
    //Residuals
    //--------------------------------------------------------------------------------------------------------------------------------

    c3->cd(1);

    histResidualsY->GetXaxis()->SetTitle("residuals y (cm)");
    histResidualsY->GetXaxis()->SetRangeUser(-0.5,0.5);
    histResidualsY->GetYaxis()->SetTitle("# of entries");
    histResidualsY->Draw();

    // Fit of Y-Residuals

    // create a TF1 with the range from -1 to 1 and 9 parameters
       TF1 *fitFcnY = new TF1("fitFcnY",fitFunction,-0.15,0.15,9);
       fitFcnY->SetNpx(10000);
       fitFcnY->SetLineWidth(1);
       fitFcnY->SetLineColor(kMagenta);

       //set parameters for fit function
       //0: G0 maximum, 1: G0 mean, 2: G0 sigma, 3: G1 maximum, 4: G1 mean, 5: G1 sigma, 6: G2 maximum, 7: G2 mean, 8: G2 sigma
       fitFcnY->SetParameter(0,10000);
       fitFcnY->SetParLimits(0,5000,18000);
       fitFcnY->SetParameter(1,0);
       fitFcnY->SetParLimits(1,-0.005,0.005);
       fitFcnY->SetParameter(2,0.0007);
       fitFcnY->SetParLimits(2,0,0.002);

       fitFcnY->SetParameter(3,25000);
       fitFcnY->SetParLimits(3,20000,60000);
       fitFcnY->SetParameter(4,0);
       fitFcnY->SetParLimits(4,-0.03,0.03);
       fitFcnY->SetParameter(5,0.02);
       fitFcnY->SetParLimits(5,0,0.03);

       fitFcnY->SetParameter(6,10000);
       fitFcnY->SetParLimits(6,5000,20000);
       fitFcnY->SetParameter(7,0);
       fitFcnY->SetParLimits(7,-0.1,0.1);
       fitFcnY->SetParameter(8,0.1);
       fitFcnY->SetParLimits(8,0,0.2);

       fitFcnY->SetParNames("G1 int","G1 mean", "G1 #sigma","G2 int","G2 mean", "G2 #sigma","G3 int","G3 mean", "G3 #sigma" );

       histResidualsY->Fit("fitFcnY","V+","ep",-0.15,0.15);


       // Draw Fit Functions:
       TF1 *Gauss_0Y = new TF1("Gauss_0Y",Gauss_0,-1,1,3);
       Gauss_0Y->SetLineColor(kGreen);
       Gauss_0Y->SetLineWidth(1);
       Gauss_0Y->SetNpx(10000);

       TF1 *Gauss_1Y = new TF1("Gauss_1Y",Gauss_1,-1,1,3);
       Gauss_1Y->SetLineColor(kRed);
       Gauss_1Y->SetNpx(10000);
       Gauss_1Y->SetLineWidth(1);

       TF1 *Gauss_2Y = new TF1("Gauss_2Y",Gauss_2,-1,1,3);
       Gauss_2Y->SetLineColor(kBlue);
       Gauss_2Y->SetNpx(10000);
       Gauss_2Y->SetLineWidth(1);

       Double_t parY[9];

       // writes the fit results into the par array
       fitFcnY->GetParameters(parY);

       Gauss_0Y->SetParameters(&parY[0]);
       Gauss_0Y->Draw("same");

       Gauss_1Y->SetParameters(&parY[3]);
       Gauss_1Y->Draw("same");

       Gauss_2Y->SetParameters(&parY[6]);
       Gauss_2Y->Draw("same");

       // draw the legend
       TLegend *legendy=new TLegend(0.6,0.45,0.88,0.25);
       legendy->SetTextFont(72);
       legendy->SetTextSize(0.04);
       legendy->AddEntry(Gauss_0Y,"Narrow Peak Fit (G1)","l");
       legendy->AddEntry(Gauss_1Y,"Residual Fit (G2)","l");
       legendy->AddEntry(Gauss_2Y,"Background Fit (G3)","l");
       legendy->AddEntry(fitFcnY,"Global Fit","l");
       legendy->Draw();

    histResidualsY->Write();
    c3->cd(2);

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

    histResidualsZ->GetXaxis()->SetTitle("residuals z (cm)");
    histResidualsZ->GetXaxis()->SetRangeUser(-0.5,0.5);
    histResidualsZ->GetYaxis()->SetTitle("# of entries");
    histResidualsZ->Draw();

    // Fit of Z-Residuals

    // create a TF1 with the range from -1 to 1 and 9 parameters
    TF1 *fitFcnZ = new TF1("fitFcnZ",fitFunction,-0.35,0.35,9);
    fitFcnZ->SetNpx(10000);
    fitFcnZ->SetLineWidth(1);
    fitFcnZ->SetLineColor(kMagenta);

    //set parameters for fit function
    //0: G0 maximum, 1: G0 mean, 2: G0 sigma, 3: G1 maximum, 4: G1 mean, 5: G1 sigma, 6: G2 maximum, 7: G2 mean, 8: G2 sigma

    //the next two parameter sets were used for either run 682 or 694

    //run 681
//    fitFcnZ->SetParameter(0,6000);
//    fitFcnZ->SetParLimits(0,5000,7000);
//    fitFcnZ->SetParameter(1,-0.1);
//    fitFcnZ->SetParLimits(1,-0.2,0);
//    fitFcnZ->SetParameter(2,0.01);
//    fitFcnZ->SetParLimits(2,0,0.1);

    //run 694
    fitFcnZ->SetParameter(0,800);
    fitFcnZ->SetParLimits(0,500,1000);
    fitFcnZ->SetParameter(1,0.25);
    fitFcnZ->SetParLimits(1,0.1,0.4);
    fitFcnZ->SetParameter(2,0.1);
    fitFcnZ->SetParLimits(2,0,1);


    fitFcnZ->SetParameter(3,16000);
    fitFcnZ->SetParLimits(3,0,40000);
    fitFcnZ->SetParameter(4,0);
    fitFcnZ->SetParLimits(4,-0.05,0.05);
    fitFcnZ->SetParameter(5,0.01);
    fitFcnZ->SetParLimits(5,0,0.1);

    fitFcnZ->SetParameter(6,3000);
    fitFcnZ->SetParLimits(6,0,10000);
    fitFcnZ->SetParameter(7,0);
    fitFcnZ->SetParLimits(7,-1,1);
    fitFcnZ->SetParameter(8,0.2);
    fitFcnZ->SetParLimits(8,0,0.6);

    fitFcnZ->SetParNames("G1 int","G1 mean", "G1 #sigma","G2 int","G2 mean", "G2 #sigma","G3 int","G3 mean", "G3 #sigma" );

    histResidualsZ->Fit("fitFcnZ","V+","ep",-0.35,0.35);

    // Draw Fit Functions:
    TF1 *Gauss_0Z = new TF1("Gauss_0Z",Gauss_0,-1,1,3);
    Gauss_0Z->SetLineColor(kGreen);
    Gauss_0Z->SetLineWidth(1);
    Gauss_0Z->SetNpx(10000);
    TF1 *Gauss_1Z = new TF1("Gauss_1Z",Gauss_1,-1,1,3);
    Gauss_1Z->SetLineColor(kRed);
    Gauss_1Z->SetLineWidth(1);
    Gauss_1Z->SetNpx(10000);
    TF1 *Gauss_2Z = new TF1("Gauss_2Z",Gauss_2,-1,1,3);
    Gauss_2Z->SetNpx(10000);
    Gauss_2Z->SetLineWidth(1);
    Gauss_2Z->SetLineColor(kBlue);

    Double_t parZ[9];

    // writes the fit results into the par array
    fitFcnZ->GetParameters(parZ);

    Gauss_0Z->SetParameters(&parZ[0]);
    Gauss_0Z->Draw("same");

    Gauss_1Z->SetParameters(&parZ[3]);
    Gauss_1Z->Draw("same");

    Gauss_2Z->SetParameters(&parZ[6]);
    Gauss_2Z->Draw("same");

    // draw the legend
    TLegend *legendz=new TLegend(0.6,0.45,0.88,0.25);
    legendz->SetTextFont(72);
    legendz->SetTextSize(0.04);
    legendz->AddEntry(Gauss_0Z,"Second Peak Fit (G1)","l");
    legendz->AddEntry(Gauss_1Z,"Residual Fit (G2)","l");
    legendz->AddEntry(Gauss_2Z,"Background Fit (G3)","l");
    legendz->AddEntry(fitFcnZ,"Global Fit","l");
    legendz->Draw();


    histResidualsZ->Write();
    c3->cd(3);
    gPad->SetLogy();
    histResidualsY->Draw();
    c3->cd(4);
    gPad->SetLogy();
    histResidualsZ->Draw();
    c3->Write();

    c4->cd(1);
    histChi2perNDF->GetXaxis()->SetTitle("chi2/NDF");
    histChi2perNDF->GetYaxis()->SetTitle("# of entries");
    histChi2perNDF->Draw();
    histChi2perNDF->Write();
    c4->cd(2);
    histPVal->GetXaxis()->SetTitle("p value");
    histPVal->GetYaxis()->SetTitle("# of entries");
    histPVal->Draw();
    histPVal->Write();
    c4->Write();

    c5->cd(1);
    histtrkTheta->GetXaxis()->SetTitle("#theta (rad)");
    histtrkTheta->GetYaxis()->SetTitle("# of entries");
    histtrkTheta->Draw();
    histtrkTheta->Write();
    c5->cd(2);
    histtrkPhi->GetXaxis()->SetTitle("#phi (rad)");
    histtrkPhi->GetYaxis()->SetTitle("# of entries");
    histtrkPhi->Draw();
    histtrkPhi->Write();
    c5->cd(3);
    histtrkMom->GetXaxis()->SetTitle("momentum (GeV/c)");
    histtrkMom->GetYaxis()->SetTitle("# of entries");
    histtrkMom->Draw();
    histtrkMom->Write();
    c5->cd(4);
    histtrkPosx->GetXaxis()->SetTitle("x-position (cm)");
    histtrkPosx->GetYaxis()->SetTitle("# of entries");
    histtrkPosx->Draw();
    histtrkPosx->Write();
    c5->cd(5);
    histtrkPosy->GetXaxis()->SetTitle("y-position (cm)");
    histtrkPosy->GetYaxis()->SetTitle("# of entries");
    histtrkPosy->Draw();
    histtrkPosy->Write();
    c5->cd(6);
    histtrkPosz->GetXaxis()->SetTitle("z-position (cm)");
    histtrkPosz->GetYaxis()->SetTitle("# of entries");
    histtrkPosz->Draw();
    histtrkPosz->Write();
    c5->Write();

    c6->cd();
    histBeamProfile->GetXaxis()->SetTitle("y-position (cm)");
    histBeamProfile->GetYaxis()->SetTitle("z-position (cm)");
    histBeamProfile->Draw("colz");
    histBeamProfile->Write();
    c6->Write();


    savefile->Close();
}