#include "TCanvas.h"
#include "TFile.h"
#include "TH1F.h"
#include "THStack.h"
#include "TLegend.h"
#include "TLegendEntry.h"
#include "TObjArray.h"
#include "TObjString.h"
#include "TString.h"
#include "TSpline.h"
#include "TSystem.h"

#include <iostream>
#include <vector>
using namespace std;

//#include "pluto_qa_thermal.C"

const Int_t nPlutoReacs = 18;

class PlutoReac : public TObject {
public:
    TString mother;
    TString dalitz;
    TString dilep;

    PlutoReac(TString m, TString dal, TString dil) {
	mother = m;
	dalitz = dal;
	dilep  = dil;
    }
};

TObjArray plutoReacs;
TObjArray plutoDilepComp;

void fillPlutoReacs() {
    plutoReacs.Add(new PlutoReac("delta", "p",   "epem"));
    //plutoReacs.Add(new PlutoReac("delta", "p",   "mpmm"));
    plutoReacs.Add(new PlutoReac("eta",   "g",   "epem"));
    plutoReacs.Add(new PlutoReac("eta",   "g",   "mpmm"));
    plutoReacs.Add(new PlutoReac("eta",   "",    "epem"));
    plutoReacs.Add(new PlutoReac("eta",   "",    "mpmm"));
    plutoReacs.Add(new PlutoReac("etap",  "g",   "epem"));
    plutoReacs.Add(new PlutoReac("etap",  "g",   "mpmm"));
    //plutoReacs.Add(new PlutoReac("jpsi",  "g",   "epem"));
    //plutoReacs.Add(new PlutoReac("jpsi",  "",    "epem"));
    //plutoReacs.Add(new PlutoReac("jpsi",  "",    "mpmm"));
    plutoReacs.Add(new PlutoReac("omega", "pi0", "epem"));
    plutoReacs.Add(new PlutoReac("omega", "pi0", "mpmm"));
    plutoReacs.Add(new PlutoReac("omega", "",    "epem"));
    plutoReacs.Add(new PlutoReac("omega", "",    "mpmm"));
    plutoReacs.Add(new PlutoReac("phi",   "",    "epem"));
    plutoReacs.Add(new PlutoReac("phi",   "",    "mpmm"));
    plutoReacs.Add(new PlutoReac("pi0",   "g",   "epem"));
    plutoReacs.Add(new PlutoReac("rho0",  "",    "epem"));
    plutoReacs.Add(new PlutoReac("rho0",  "",    "mpmm"));

    plutoDilepComp.Add(new PlutoReac("eta",   "g",   "epem"));
    plutoDilepComp.Add(new PlutoReac("eta",   "",    "epem"));
    //plutoDilepComp.Add(new PlutoReac("jpsi",  "",    "epem"));
    plutoDilepComp.Add(new PlutoReac("omega", "pi0", "epem"));
    plutoDilepComp.Add(new PlutoReac("omega", "",    "epem"));
    plutoDilepComp.Add(new PlutoReac("phi",   "",    "epem"));
    plutoDilepComp.Add(new PlutoReac("rho0",  "",    "epem"));
}

TString getDecayName(const TString &name) {
    TString mother("");
    TString dalitz("");
    TString dilep("");

    if(name.Contains("Delta+."))
	mother = "#Delta^{+}";
    if(name.Contains("eta."))
        mother = "#eta";
    if(name.Contains("etap."))
        mother = "#eta '";
    if(name.Contains("omega."))
        mother = "#omega";
    if(name.Contains("pi0.") && !name.Contains("omega."))
        mother = "#pi^{0}";
    if(name.Contains("phi."))
        mother = "#phi";
    if(name.Contains("rho0."))
        mother = "#rho^{0}";
    if(name.Contains("jpsi."))
	mother = "J/#Psi";
    if(name.Contains(".pepem") || name.Contains(".pmpmm"))
        dalitz = "p";
    if(name.Contains(".pi0"))
        dalitz = "#pi^{0}";
    if(name.Contains(".g"))
        dalitz = "#gamma";
    if(name.Contains("epem"))
        dilep = "e^{+} e^{-}";
    if(name.Contains("mpmm"))
        dilep = "#mu^{+} #mu^{-}";

    return TString(mother + "#rightarrow" + dalitz + dilep);
}

Color_t getDecayColor(const TString &name) {

    if(name.Contains("Delta+"))
        return kBlack;
    if(name.Contains("eta.")) {
	if(name.Contains(".epem") || name.Contains(".mpmm")) {
	    return kBlue;
	}
	if(name.Contains(".gepem") || name.Contains(".gmpmm")) {
            return kBlue - 10;
	}
    }
    if(name.Contains("etap.")) {
        return kAzure +10;
    }
    if(name.Contains("omega.")) {
	if(name.Contains(".pi0epem") || name.Contains(".pi0mpmm")) {
            return kSpring+10;
	} else {
	    return kGreen;
	}
    }
    if(name.Contains("pi0."))
        return kRed;
    if(name.Contains("phi."))
        return kCyan;
    if(name.Contains("rho0."))
        return kMagenta;
    if(name.Contains("jpsi.")) {
	if(name.Contains(".gepem") || name.Contains(".gmpmm")) {
	    return kOrange-5;
	} else {
	    return kOrange;
	}
    }
    cout<<"No color found for "<<name.Data()<<endl;
    return kGray;
}

void mergeInvMassHists(TObjArray const* const inFilesList,
                       THStack *hinvmassStack, TLegend *legend,
                       TObjArray *splines) {

    const Int_t nentries = inFilesList->GetEntries();
    //TH1F  *hinvmass[nentries];
    //TFile *f[nentries];

    TH1F  *hinvmass[20];
    TFile *f[20];

    splines->Clear();

    for (Int_t i = 0; i < nentries; i++) {
        TString fileName = ((TObjString*)inFilesList->At(i))->GetString();
        f[i] = new TFile(fileName.Data());
        hinvmass[i] = (TH1F*)gDirectory->Get("h_invmass;1");

        hinvmass[i]->SetLineColor(getDecayColor(fileName));
        hinvmass[i]->SetLineWidth(2);

        if(fileName.Contains("epem")) {
            hinvmass[i]->GetXaxis()->SetTitle("M_{ee} [GeV/c^{2}]");
            hinvmass[i]->GetYaxis()->SetTitle("dN/dM_{ee} [1/(GeV/c^{2})]");
            hinvmass[i]->GetYaxis()->SetTitleOffset(1.3);
            hinvmass[i]->GetYaxis()->SetRangeUser(1.e-4, 1.e2);
            //hinvmass[i]->SetTitle("Pluto dielectron cocktail");
        }

        if(fileName.Contains("mpmm")) {
            hinvmass[i]->GetXaxis()->SetTitle("M_{#mu#mu} [GeV/c^{2}]");
            hinvmass[i]->GetYaxis()->SetTitle("dN/dM_{#mu#mu} [1/(GeV/c^{2})]");
            hinvmass[i]->GetYaxis()->SetTitleOffset(1.3);
            hinvmass[i]->GetYaxis()->SetRangeUser(1.e-4, 1.e0);
            //hinvmass[i]->SetTitle("Pluto dimuon cocktail");
        }

        TString decName = getDecayName(fileName);
        legend->AddEntry(hinvmass[i], decName.Data(), "lpf");

        hinvmassStack->Add(hinvmass[i]);

        TSpline3 *sp = new TSpline3(hinvmass[i]);
        sp->SetLineColor(getDecayColor(fileName));
        sp->SetLineWidth(2);
        splines->Add(sp);
    }
}

void showInvMass(const TString &inDir, const TString &inFilePattern,
		 Int_t energy, Bool_t dimuon) {

    TString inFiles = gSystem->GetFromPipe("ls -1 " + inDir + "/" + inFilePattern);

    TObjArray *inFilesList = inFiles.Tokenize("\n");

    if(inFilesList->GetEntries() == 0) {
	cout<<"File list is empty."<<endl;
	return;
    }

    TString lplus    = "e+";
    TString lminus   = "e-";
    TString ll       = "ee";
    TString dilep    = "dielectron";

    if(dimuon) {
	lplus    = "mu+";
	lminus   = "mu-";
	ll       = "#mu#mu";
        dilep    = "dimuon";
    }

    THStack   *hinvmassStack = new THStack("hinvmstack_" + inFilePattern,"Di-Lepton mass");
    TLegend   *legend        = new TLegend(0.69, 0.88 - (0.3/7.) * inFilesList->GetEntries(), 0.89, 0.88, "","brNDC");
    TObjArray *splines       = new TObjArray(inFilesList->GetEntries());

    TString histTitle(Form("Pluto %s cocktail at %i GeV", dilep.Data(), energy));

    mergeInvMassHists(inFilesList, hinvmassStack, legend, splines);

    TCanvas *c_invmass = new TCanvas("c_invmass_" + inFilePattern,histTitle.Data(),10,10,600,600);

    c_invmass->cd(1)->SetLogy();

    hinvmassStack->Draw("hist l nostack");
    hinvmassStack->GetXaxis()->SetTitle("M_{"+ll+"} [GeV/c^{2}]");
    hinvmassStack->GetYaxis()->SetTitle("dN/dM_{"+ll+"} [1/(GeV/c^{2})]");
    hinvmassStack->GetYaxis()->SetTitleOffset(1.3);
    hinvmassStack->GetYaxis()->SetRangeUser(0.001, 3000.);
    hinvmassStack->SetTitle(histTitle);

    legend->Draw("same");
}

TH1F* get1dHist(TLegend * const legend,
		const TString &inDir, Int_t energy,
		const TString &part, const TString &dalitz,
		const TString &lplm, const TString &hist) {

    TString fileName  = Form("qa.pluto.%i",energy);
    fileName = fileName + "gev." + part + "." + dalitz + lplm + ".root";

    TFile::Open(inDir + "/" + fileName);
    TH1F *h = (TH1F*)gDirectory->Get(hist.Data());

    if(h) {
	Color_t col = fileName.Contains("epem")?kRed:kBlue;
	h->SetLineColor(col);
        TString decName = getDecayName(fileName);
        legend->AddEntry(h, decName.Data(), "lpf");
    }
    else
        cout<<"Histogram "<<hist.Data()<<" not found in file "<<inDir.Data()<<"/"<<fileName.Data()<<endl;

    return h;
}

void showDist(const TString &inDir, Int_t energy, const TString &h_dist) {

    Int_t nHists = plutoDilepComp.GetEntries();
    TString histName = h_dist + ";1";

    TCanvas *c_p_25gev = new TCanvas(Form("c_p_%igev_%s", energy, histName.Data()),
				     Form("Overlay %s of dielectron/dimuon channels at %i GeV",
					  histName.Data(), energy),
				     10, 10, 900, 600);
    c_p_25gev->Divide(3,2);

    TLegend *legends[nPlutoReacs];
    TH1F *h_epem[nPlutoReacs];
    TH1F *h_mpmm[nPlutoReacs];


    for(Int_t i = 0; i < nHists; i++) {

	c_p_25gev->cd(i+1);

	legends[i] = new TLegend(0.69, 0.65, 0.89, 0.75, "","brNDC");

        PlutoReac *pluto = (PlutoReac*)plutoDilepComp.At(i);

	h_epem[i] = get1dHist(legends[i], inDir, energy, pluto->mother, pluto->dalitz,"epem", histName.Data());
	if(h_epem[i]) {
	    h_epem[i]->Draw("");
	}
	h_mpmm[i] = get1dHist(legends[i], inDir, energy, pluto->mother, pluto->dalitz, "mpmm", histName.Data());
	if(h_mpmm[i]) {
	    h_mpmm[i]->Draw("same");
	}
	legends[i]->Draw("same");
    }


}

void plot_pluto_cocktail(TString inDir = "./cktA") {

    gROOT->SetStyle("Default");

    fillPlutoReacs();

    const Int_t nEb = 2;
    Int_t arrEb[nEb] = { 8, 25 };

    const Int_t nDists = 4;
    TString arrDists[nDists] = { "h_p", "h_pt", "h_y", "h_invmass" };

    TString inFilePrefix = "qa.pluto";
    TString inFilePattern("*root");

    for(Int_t e = 0; e < nEb; e++) {
        Int_t Eb = arrEb[e];
	TString sEb     = Form("%igev", Eb);

        inFilePattern = inFilePrefix + "." + sEb + "*" + "epem" + ".root";
	showInvMass(inDir, inFilePattern, Eb, kFALSE);

        inFilePattern = inFilePrefix + "." + sEb + "*" + "mpmm" + ".root";
	showInvMass(inDir, inFilePattern, Eb, kTRUE);

	for(Int_t d = 0; d < nDists; d++) {
	    showDist(inDir, Eb, arrDists[d]);
	}
    }
}