#include <iostream>
#include <vector>

#include "TFile.h"
#include "TChain.h"
#include "TTree.h"
#include "TFileCollection.h"
#include "TCollection.h"
#include "THashList.h"
#include "TStopwatch.h"
#include "TMath.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TVector3.h"
#include "TString.h"

#include "Event.h"
#include "Event.cxx"
#include "myAnaConsts.h"

using namespace std;

void corr2D(std::string datalist)
{
   std::cout << datalist << std::endl;

   TFile fout(TString::Format("fout%s.root", datalist.c_str()), "recreate");

   TChain *chain_d0 = new TChain("npd0ana1/VertexCompositeNtuple");
   TChain *chain_tracks = new TChain("track_ana/trackTree");

   TFileCollection* fcData = new TFileCollection(datalist.c_str(), "", datalist.c_str());

   chain_d0->AddFileInfoList(fcData->GetList());
   std::cout << "d0 ready" << std::endl;

   chain_tracks->AddFileInfoList(fcData->GetList());
   std::cout << "tracks ready" << std::endl;

   Event* evt = new Event(chain_d0, chain_tracks);

   // declare hists
   //
   map<string, TH2*> hSignal_D0[ana::nMass][ana::nPt];

   for(int ipt=0; ipt<ana::nPt; ipt++){
      for(int imass=0; imass<ana::nMass; imass++){
         (hSignal_D0[imass][ipt])["largeDCA"] = new TH2D(Form("hSignal_mass%d_pt%d_largeDCA", imass, ipt),
               "", ana::nEtaBin, ana::etaBegin, ana::etaEnd,
               32, -1.57, 4.71);
         (hSignal_D0[imass][ipt])["smallDCA"] = new TH2D(Form("hSignal_mass%d_pt%d_smallDCA", imass, ipt),
               "", ana::nEtaBin, ana::etaBegin, ana::etaEnd,
               32, -1.57, 4.71);
      }
   }


   for(int imass=0; imass<ana::nMass; imass++){
      for(int ipt=0; ipt<ana::nPt; ipt++){
         for(auto& h : hSignal_D0[imass][ipt]) {
            h.second->SetCanExtend(0);
            std::cout << "xmax of y axis" << h.second->GetYaxis()->GetXmax() << std::endl;;
         }
      }
   }

   // declare vectors
   vector<TVector3> pVect_trg_d0[ana::nMass][ana::nPt];

   vector<TVector3> pVect_dau1_d0[ana::nMass][ana::nPt];
   vector<TVector3> pVect_dau2_d0[ana::nMass][ana::nPt];

   vector<TVector3> pVect_ass;

   // start timing
   TStopwatch ts;
   ts.Start();

   // loop
   //
   // temporary vectors
   TVector3 p_dau1(0, 0, 0), p_dau2(0, 0, 0), p_d0(0, 0, 0), p_ass(0, 0, 0);

   std::cout << evt->GetEntries() << std::endl;
   long int nentries = evt->GetEntries();
   int percent = 0;
   long int skip = 0;
   for(long int ientry=0; ientry<nentries; ientry++){
      int current = (int)ientry*100/nentries;
      if( current == 5*percent){
         std::cout << 5*percent++ << "\% percents completed" << std::endl;
      }
      auto bytes = evt->GetEntry(ientry);
      if(bytes == 0 || bytes == -1) {
         skip++;
         continue;
      }

      for(int id0=0; id0<evt->CandSize(); id0++){
         int imass = ana::findMassBin(evt->Mass(id0));
         int ipt = ana::findPtBin(evt->Pt(id0));
         int iy = ana::findYBin(evt->Y(id0));
         if(imass == -1) continue;
         if(ipt == -1) continue;
         if(iy == -1) continue;

         p_dau1.SetPtEtaPhi(evt->PtD1(id0), evt->etaD1(id0), evt->phiD1(id0));
         p_dau2.SetPtEtaPhi(evt->PtD2(id0), evt->etaD2(id0), evt->phiD2(id0));
         p_d0 = p_dau1 + p_dau2;

         pVect_trg_d0[imass][ipt].push_back(p_d0);
         pVect_dau1_d0[imass][ipt].push_back(p_dau1);
         pVect_dau2_d0[imass][ipt].push_back(p_dau2);
      }

      for(unsigned int itrack=0; itrack<evt->CandSizeTrk(); itrack++){
         p_ass.SetPtEtaPhi(evt->PtTrk(itrack), evt->EtaTrk(itrack), evt->PhiTrk(itrack));
         pVect_ass.push_back(p_ass);
      }

      // calculate signal
      unsigned int nMult_ass = (unsigned int) pVect_ass.size();
      
      unsigned int nMult_trg_raw_d0[ana::nMass][ana::nPt] = {0}; // Ntrig for mass & pt bins
      double nMult_trg_eff_d0[ana::nMass][ana::nPt] = {0.}; // eff corrected Ntrig for mass & pt bins

      for(int imass=0; imass<ana::nMass; imass++){
         for(int ipt=0; ipt<ana::nPt; ipt++){
            unsigned int nMult_trg_d0 = (unsigned int) pVect_trg_d0[imass][ipt].size();
            for(unsigned int id0=0; id0<nMult_trg_d0; id0++){
               double effks = 1.0;
               nMult_trg_raw_d0[imass][ipt] += 1;
               nMult_trg_eff_d0[imass][ipt] += 1./effks;
            }

            for(unsigned int id0=0; id0<nMult_trg_d0; id0++){
               if(ipt!=ana::findPtBin(pVect_trg_d0[imass][ipt].at(id0).Pt())) std::cout << "pT bin error" << std::endl;
               double effks = 1.0;
               for(unsigned iass=0; iass<nMult_ass; iass++){
                  TVector3 pvector_ass = pVect_ass.at(iass);
                  double effweight_ass = evt->WeightTrk(iass);

                  double deltaEta = pvector_ass.Eta() - pVect_trg_d0[imass][ipt].at(id0).Eta();
                  double deltaPhi = pvector_ass.DeltaPhi(pVect_trg_d0[imass][ipt].at(id0));
                  if(deltaPhi>-ana::PI && deltaPhi<-ana::PI/2.) deltaPhi += 2*ana::PI;

                  (hSignal_D0[imass][ipt])["largeDCA"]->Fill(deltaEta,
                        deltaPhi,
                        1./nMult_trg_eff_d0[imass][ipt]/effks/effweight_ass);
               }
            }
         }
      }

      if(ientry==nentries/2){
         for(int imass=0; imass<ana::nMass; imass++){
            for(int ipt=0; ipt<ana::nPt; ipt++){
               for(auto& h : hSignal_D0[imass][ipt]) {
                  std::cout << "xmax of y axis" << h.second->GetYaxis()->GetXmax() << std::endl;;
               }
            }
         }
      }

      // clear all
      p_dau1.SetXYZ(0, 0, 0);
      p_dau2.SetXYZ(0, 0, 0);
      p_d0.SetXYZ(0, 0, 0);
      p_ass.SetXYZ(0, 0, 0);
      for(int imass=0; imass<ana::nMass; imass++){
         for(int ipt=0; ipt<ana::nPt; ipt++){
            pVect_trg_d0[imass][ipt].clear();
            pVect_dau1_d0[imass][ipt].clear();
            pVect_dau2_d0[imass][ipt].clear();
         }
      }
      pVect_ass.clear();
      pVect_ass.clear();
   }
   std::cout << "completed loop" << std::endl;
   std::cout << skip << " events are skipped" << std::endl;

   ts.Stop();
   ts.Print();

   fout.cd();
   
   for(int ipt=0; ipt<ana::nPt; ipt++){
      for(int imass=0; imass<ana::nMass; imass++){
         for(auto& h : hSignal_D0[imass][ipt]) h.second->Write();
      }
   }

   for(int ipt=0; ipt<ana::nPt; ipt++){
      for(int imass=0; imass<ana::nMass; imass++){
         for(auto& h : hSignal_D0[imass][ipt]) {
            std::cout << "xmax of y axis" << h.second->GetYaxis()->GetXmax() << std::endl;;
            delete h.second;
         }
      }
   }

   delete evt;

   return ;
}