#include "PParticle.h"
#include "TMath.h"
#include "TLorentzVector.h"
#include "TChain.h"
#include "TTree.h"
#include "TClonesArray.h"
#include "TFile.h"
#include "TH1F.h"
#include "TCanvas.h"
#include "TMath.h"

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

void plotPluto()
{

    Bool_t reco          = kTRUE;
    Bool_t momCut        = kFALSE;
    Bool_t oaCut         = kTRUE;
    Bool_t acceptanceCut = kFALSE;
    Bool_t useWeights    = kTRUE;

    TChain data("data");
    if(reco) for(Int_t i = 1; i <= 20; i++) data.AddFile(Form("/lustre/hades/user/hadesdst/feb22/convert_leptons/pp45_%i_pluto.root",i));
    else     for(Int_t i = 1; i <= 20; i++) data.AddFile(Form("/lustre/hades/user/hadesdst/feb22/convert_leptons_noreco/pp45_%i_pluto.root",i));


   TClonesArray *evt=new TClonesArray("PParticle",100);
   data.SetBranchAddress("Particles",&evt);

   TFile* out = 0 ;

   if (reco) out = new TFile("pluto_pp45_reco.root","RECREATE");
   else      out = new TFile("pluto_pp45_noreco.root","RECREATE");
   out->cd();

   //-----------------------------------------------------------------
   // helper maps
   map<Int_t, vector<Int_t> >      mIDs;
   map<Int_t, map<Int_t,vector<PParticle*> > > mPluto;
   map<Int_t, TH1F*>               mH;
   map<Int_t, TString >            mnames;
   map<Int_t, Double_t >           mBR;

   map<Int_t, vector<Int_t> >::iterator      itIDs;
   map<Int_t, map<Int_t,vector<PParticle*> > >::iterator itPluto;
   map<Int_t, TString >::iterator            itnames;
   map<Int_t, TH1F* >::iterator              itH;
   map<Int_t, Double_t >::iterator           itBR;

   Int_t color[] = {kRed,kMagenta+1,kBlue,kCyan+1,kGreen+1,kOrange+1};

   mnames[75101 ]="pi0->e+e-g";
   mnames[175101]="eta->e+e-g";
   mnames[410203]="rho0->e+e";
   mnames[525107]="omega->e+e-pi0";
   mnames[520203]="omega->e+e-";
   mnames[550203]="phi->e+e-";

   map<Int_t,vector<PParticle*> > mInit;
   vector<PParticle*> dPluto;
   vector<Int_t> d;
   d.push_back(2);
   d.push_back(3);
   d.push_back(1);

   mIDs  [75101]  = d;   //pi0-> e+e-g
   mIDs  [175101] = d;   //et-> e+e-g
   mPluto[75101]  = mInit;
   mPluto[175101] = mInit;

   d.clear();
   d.push_back(2);
   d.push_back(3);

   mIDs  [410203] = d; //rho0-> e+e
   mIDs  [520203] = d; //w-> e+e-
   mIDs  [550203] = d; //phi-> e+e-
   mPluto[410203] = mInit;
   mPluto[520203] = mInit;
   mPluto[550203] = mInit;

   d.push_back(7);

   mIDs  [525107] = d;  // w->e+e- pi0
   mPluto[525107] = mInit;
   //-----------------------------------------------------------------


   //-----------------------------------------------------------------
   // hists + canvas
   TCanvas* result =  new TCanvas("cmass","Invariant mass",1000,800);
   Int_t ct= 0;
   Int_t nBin = 150;
   for (itnames = mnames.begin(); itnames != mnames.end(); ++itnames){
        mH[itnames->first] = new TH1F(itnames->second,itnames->second,nBin,0,1.500);
	mH[itnames->first]->SetLineColor(color[ct]);
	ct++;
   }

   mH[0] = new TH1F("Sum","Sum",nBin,0,1.500);
   mH[0] ->SetLineColor( kBlack);
   //-----------------------------------------------------------------


   //-----------------------------------------------------------------
   //-----------------------------------------------------------------
   // loop data
   Long64_t nentries = data.GetEntries();
   cout<<"ENTRIES "<<nentries<<endl;
   //nentries=100;
   Long64_t nbytes = 0;
   for (Long64_t i=0; i<nentries;i++) {
      nbytes += data.GetEntry(i);
      if( (i%100000) == 0) cout<<"------------------------Event "<<i<<endl;

      //-----------------------------------------------------------------
      // clear the decay map  for new event
      for (itPluto = mPluto.begin(); itPluto != mPluto.end(); ++itPluto) {
          if(itPluto->second.size() == 0) continue;
          itPluto->second.clear();
      }
      //-----------------------------------------------------------------

      //-----------------------------------------------------------------
      // loop event to fill decay map
      for(Int_t j = 0 ; j < evt->GetEntries(); j++){
	  PParticle* p = (PParticle*)evt->At(j);
	  Int_t id   = p->GetSourceId();
	  Int_t pidx = p->GetParentIndex();

	  itIDs = mIDs.find(id);
	  if(itIDs == mIDs.end()) continue;
          itnames = mnames.find(id);

	  itPluto = mPluto.find(id);
          map<Int_t , vector<PParticle*> >& mL = itPluto->second;
	  map<Int_t , vector<PParticle*> >::iterator iL ;

	  iL = mL.find(pidx);
	  if(iL != mL.end()) {
	      iL->second.push_back(p);
	  } else {
	      vector<PParticle*> v;
	      v.push_back(p);
              mL[pidx] = v;
	  }
      }
      //-----------------------------------------------------------------

      //-----------------------------------------------------------------
      // analyze the found decays
      for (itPluto = mPluto.begin(); itPluto != mPluto.end(); ++itPluto) {
          itnames = mnames.find(itPluto->first);
	  map<Int_t,vector<PParticle*> > & mL = itPluto->second ;
	  
	  if(mL.size() == 0) continue;

          map<Int_t , vector<PParticle*> >::iterator iL ;
          for (iL = mL.begin(); iL != mL.end(); ++iL) {
              vector<PParticle*>& dP = iL->second ;

	      //-----------------------------------------------------------------
	      // calculate mother particle
	      TLorentzVector mother;
	      PParticle* em = 0;
	      PParticle* ep = 0;
	      for(UInt_t k = 0 ; k < dP.size(); k++){
                  if(dP[k]->ID() == 1) continue; //skip neutral daughters, they would not be reconstructed
                  if(dP[k]->ID() == 7) continue;
                  if(dP[k]->ID() == 2) ep = dP[k];
                  if(dP[k]->ID() == 3) em = dP[k];
		  mother += *dP[k];
	      }
	      //-----------------------------------------------------------------

	      //-----------------------------------------------------------------
	      if(ep && em){

		  //-----------------------------------------------------------------
                  // apply selected cuts
		  Double_t oA   = ep->Angle(em->Vect()) * TMath::RadToDeg();
		  Double_t thep = ep->Theta()         * TMath::RadToDeg();
		  Double_t them = em->Theta()         * TMath::RadToDeg();

		  if(momCut){
		      if(em->P() < 0.1) continue;
		      if(ep->P() < 0.1) continue;
		  }

		  if(acceptanceCut){
		      if(thep < 20 || thep > 85 ) continue;
		      if(them < 20 || them > 85 ) continue;
		  }

		  if(oaCut){
		      if(oA < 9) continue;
		  }
		  //-----------------------------------------------------------------


		  //-----------------------------------------------------------------
                  // finally fill hists
		  Double_t weight = 1;

		  if(useWeights) {
		      weight = ep->W();  // normal branching, but for multiple channels nchan * branching
		  }

		  itH = mH.find(itPluto->first);
		  itH->second->Fill(mother.M(),weight);
		  mH[0]->Fill(mother.M(),weight); // add to the SUM hist
		  //-----------------------------------------------------------------
	      }
	      //-----------------------------------------------------------------
	  }
      }
      //-----------------------------------------------------------------

   }
   //-----------------------------------------------------------------
   //-----------------------------------------------------------------


   //-----------------------------------------------------------------
   // write hist + canvas
   out->cd();

   result->cd();
   mH[0]->DrawCopy();
   for (itH = mH.begin(); itH != mH.end(); ++itH){
       itH->second->DrawCopy("same");
       itH->second->Write();
   }

   result->Write();
   out->Close();
   //-----------------------------------------------------------------
   return;
}