#include "PParticle.h"

#include "TCanvas.h"
#include "TChain.h"
#include "TClonesArray.h"
#include "TDatabasePDG.h"
#include "TFile.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TLorentzVector.h"
#include "TObjArray.h"
#include "TString.h"
#include "TSystem.h"
#include "TStyle.h"

TDatabasePDG *pdg;
void getParticleNamePDG(TString const& fileName, TString &part) {
    // Derive PDG particle name from Pluto sim file name.

    TObjArray *tokens = fileName.Tokenize(".");
    // Index to particle name.
    Int_t iPart = tokens->GetEntries() - 1 - 3;
    if(iPart >= 0) {
	TObjString *oPart = (TObjString*)tokens->At(iPart);
        // Found particle name.
	part = oPart->String();
        // File names for Eta prime and J/Psi differ from PDG names.
	if(part == "etap") {
	    part = "eta'";
	}
	if(part == "jpsi") {
	    part = "J/psi";
	}
    } else{
	cout<<"Could not extract particle name from file "<<fileName.Data()<<endl;
        part = "";
    }
}

Int_t pluto_qa_thermal(TString inFilePattern="/hera/hades/user/ekrebs/cbm/pluto/auaubw/25gev/rho_eecutoff/epem/*root",
		       TString outFile="./histosbw/ana_pluto_25gev_rho_eecutoff_epem.root",
		       Bool_t dimuon=kFALSE, Int_t maxFiles = -1) {
    // Analyses Pluto sim files.
    // Plots invariant mass, p_t, y, p_t(y), m_t - pole mass, opening angle,
    // momentum of dileptons and individual leptons, momenta in x,y and z.
    // Histograms need to be scaled afterwards with pluto_qa_scale_hists
    // macro.
    //
    // inFilePattern: Name of Pluto sim file(s). May be a regular expression.
    // outFile:       Output file with histograms.
    // dimuon:        Specifies whether files contain dielectrons or dimuons.
    // maxFiles:      Maximum number of files to analyze. If this number is
    //                negative then all the files that match the pattern are
    //                analyzed.

    // load libraries
    //gSystem->Load("libPluto.so");

    // Used to identify PParticle.
    TString lplus    = "e+";
    TString lminus   = "e-";

    // Used in histogram titles only.
    TString ll       = "ee";

    // Analyze dimuon decay channel instead of dielectron.
    if(dimuon) {
        lplus    = "mu+";
        lminus   = "mu-";
        ll       = "#mu#mu";
    }

    gStyle->SetPalette(1);

    // Reset ROOT and connect tree file
    gROOT->Reset();

    // Get list of input files
    TString inFiles = gSystem->GetFromPipe("ls -1 " + inFilePattern);
    TObjArray *inFilesList = inFiles.Tokenize("\n");
    if(inFilesList->GetEntries() == 0) {
	cout<<"Input files: "<<inFilePattern.Data()<<endl;
        return 1;
    }

    TChain *Reaction = new TChain("data");
    Int_t i = 0;
    Int_t nFiles = inFilesList->GetEntries();
    if(maxFiles >= 0) nFiles = TMath::Min(inFilesList->GetEntries(), maxFiles);
    while(i < nFiles) {
        TString inFile = ((TObjString*)inFilesList->At(i))->GetString();
	Reaction->Add(inFile.Data());
	if(i < 20 || i == nFiles-1) {
	    cout<<"Adding file "<<inFile.Data()<<endl;
	}
	if(i == 20) {
	    cout<<"..."<<endl;
	}
        i++;
    }

    TClonesArray *evt=new TClonesArray("PParticle",100);    // must create TClonesArray
    Reaction->SetBranchAddress("Particles",&evt);         // recover PParticles

    // Get particle mass.
    TString part;
    getParticleNamePDG(((TObjString*)inFilesList->At(0))->String(), part);
    pdg = new TDatabasePDG();
    TParticlePDG *partPDG = pdg->GetParticle(part.Data());
    Float_t m0 = 0.F;
    if(partPDG) {
	m0 = partPDG->Mass();
    } else{
	cout<<"Particle with name "<<part.Data()<<" not found in PDG database."<<endl;
    }

    // Define histograms
    Float_t maxM = 3.3F;

    // Invariant mass
    TH1F *h_invmass = new TH1F("h_invmass","Di-Lepton mass",100*maxM, 0., maxM);
    h_invmass->GetXaxis()->SetTitle("M_{"+ll+"} #left[ #frac{GeV}{c^{2}} #right]");
    h_invmass->GetYaxis()->SetTitle("dN/dM_{"+ll+"} #left[ #left( #frac{GeV}{c^{2}} #right)^{-1} #right]");
    h_invmass->GetYaxis()->SetTitleOffset(1.3);
    h_invmass->SetLineWidth(2);

    // Opening angle in degrees
    TH1F *h_opang = new TH1F("h_opang", "Opening angle [#circ]",180,0,180);
    h_opang->GetXaxis()->SetTitle("#Theta_{"+ll+"} [#circ]");
    h_opang->GetYaxis()->SetTitle("#frac{dN}{d#Theta_{"+ll+"}} #left[ #frac{1}{#circ} #right]");

    // Rapidity over transverse momentum
    TH2F* h_pty = new TH2F("h_pty","",500,-1.,4.,75,0,1.5);
    h_pty->GetXaxis()->SetTitle("y");
    h_pty->GetYaxis()->SetTitle("p_{t} #left[ #frac{MeV}{c} #right]");
    h_pty->GetYaxis()->SetTitleOffset(1.3);

    // Rapidity distribution
    TH1F* h_y = new TH1F("h_y","Rapidity",500,-1.,4.);
    h_y->GetXaxis()->SetTitle("y");
    h_y->GetYaxis()->SetTitle("#frac{dN}{dY}");

    // Momentum of the dileptons
    TH1F *h_p = new TH1F("h_p", "Dilepton Momentum Distribution", 500, 0., 15);
    h_p->GetXaxis()->SetTitle("p #left[ #frac{MeV}{c} #right]");

    // Transverse momentum of the dileptons
    TH1F *h_pt    = new TH1F("h_pt",    "Dilepton Transverse Momentum", 200, 0, 2);
    h_pt->GetXaxis()->SetTitle("p_{T} #left[ #frac{GeV}{c} #right]");
    h_pt->GetYaxis()->SetTitle("#frac{dN}{dp_{T}} #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    // Transverse mass of the dileptons
    TH1F *h_mt    = new TH1F("h_mt",    "Dilepton Transverse Mass", 200, 0, 2);
    h_mt->GetXaxis()->SetTitle("m_{T} #left[ #frac{GeV}{c} #right]");
    h_mt->GetYaxis()->SetTitle("#frac{1}{m_{T}} #frac{dN}{dm_{T}} #left[ #left( #frac{GeV}{c^3} #right)^{-1} #right]");

    // Transverse mass of the dileptons minus pole mass
    TH1F *h_mt_m0 = new TH1F("h_mt_m0",    "Dilepton Transverse Mass Minus Pole Mass", 200, 0, 2);
    h_mt_m0->GetXaxis()->SetTitle("m_{T} - <m_{0}> #left[ #frac{GeV}{c} #right]");
    h_mt_m0->GetYaxis()->SetTitle("#frac{1}{m_{T}} #frac{d^{2}N}{dm_{T} dy} #left[ #left( #frac{GeV}{c^{3}} #right)^{-1} #right]");
    h_mt_m0->GetYaxis()->SetTitleOffset(1.3);

    // Momentum of positive lepton
    TH1F *h_p1   = new TH1F("h_p1", "Momentum of "+lplus, 500,0,15);
    h_p1->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_p1->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");
    //fh_p1->SetLineWidth(2);

    // Momentum of negative leptons
    TH1F *h_p2   = new TH1F("h_p2", "Momentum of "+lminus, 500,0,15);
    h_p2->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_p2->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    // Momenta in x,y,z of individual leptons
    TH1F *h_px     = new TH1F("h_px", "Momentum x", 500,0,15);
    h_px->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_px->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    TH1F *h_py     = new TH1F("h_py", "Momentum y", 500,0,15);
    h_py->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_py->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    TH1F *h_pz     = new TH1F("h_pz", "Momentum z", 500,0,15);
    h_pz->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_pz->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    // Momenta in x,y,z of dileptons
    TH1F *h_px_ll     = new TH1F("h_px_ll", "Dilepton Momentum x", 500,0,15);
    h_px_ll->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_px_ll->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    TH1F *h_py_ll     = new TH1F("h_py_ll", "Dilepton Momentum y", 500,0,15);
    h_py_ll->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_py_ll->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    TH1F *h_pz_ll     = new TH1F("h_pz_ll", "Dilepton Momentum z", 500,0,15);
    h_pz_ll->GetXaxis()->SetTitle("p #left[ #frac{GeV}{c} #right]");
    h_pz_ll->GetYaxis()->SetTitle("dN/dp #left[ #left( #frac{GeV}{c} #right)^{-1} #right]");

    // Loop over Pluto files.
    PParticle *ep,*em = NULL;

    Int_t nEntries = Reaction->GetEntries();

    for (Int_t i = 0; i < nEntries; i++) {
	if(i % 100000 == 0) cout<<"Processing entry "<<i<<endl;
	ep = NULL;
	em = NULL;
	Reaction->GetEntry(i);

	for (int j = 0;j < evt->GetEntriesFast(); j++) {
	    PParticle *current =(PParticle*)evt->At(j);
	    if (current->Is(lplus.Data())) ep=current;
	    if (current->Is(lminus.Data())) em=current;
	}
	
	if (ep && em) {
	    //particles found

	    //It is very important to parse the PParticle
	    //to a TLorentzVector, because the "+" operator
	    //is reserved for a reaction (i.e. a compound
	    //particle is created

	    TLorentzVector  dilepton =
		(*(TLorentzVector *) ep) + (*(TLorentzVector *) em);
	    h_invmass->Fill(dilepton.M());

	    Double_t pt  = dilepton.Pt();
            Double_t mt  = dilepton.Mt();
            Double_t rap = dilepton.Rapidity();

            Float_t oAngle=ep->Angle(em->Vect())*TMath::RadToDeg();

	    h_pty->Fill(rap,pt);
            h_y->Fill(rap);

            h_p->Fill(dilepton.P());

	    h_pt->Fill(pt);
            h_mt->Fill(mt);
            h_mt_m0->Fill(mt - m0, 1./(mt));

            h_p1->Fill(ep->P());
            h_p2->Fill(em->P());

            h_px->Fill(ep->Px());
            h_py->Fill(ep->Py());
            h_pz->Fill(ep->Pz());

            h_px->Fill(em->Px());
            h_py->Fill(em->Py());
            h_pz->Fill(em->Pz());

            h_px_ll->Fill(dilepton.Px());
            h_py_ll->Fill(dilepton.Py());
            h_pz_ll->Fill(dilepton.Pz());

            h_opang->Fill(oAngle);
	}
    }

    // Write histograms.
    TFile fout(outFile.Data(),"RECREATE");
    fout.cd();
    h_invmass->Write();
    h_opang->Write();
    h_pty->Write();
    h_y->Write();
    h_p->Write();
    h_pt->Write();
    h_mt->Write();
    h_mt_m0->Write();
    h_p1->Write();
    h_p2->Write();
    h_px->Write();
    h_py->Write();
    h_pz->Write();
    h_px_ll->Write();
    h_py_ll->Write();
    h_pz_ll->Write();
    fout.Close();

    delete pdg;

    return 0;

}