// ==============================================================================
//  File: getv2.C
//
//
//  Author: Thomas Ullrich
//  Last update: March, 2016
//
// ==============================================================================
#include <iostream>
#include "TObject.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TTree.h"
#include "TMath.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TStyle.h"
#include "TChain.h"
#include "TLorentzVector.h"
#include "TText.h"
#include <fstream>
#include <string>
#define PR(x) cout << " = " << (x) << endl;

using namespace std;

#define NEW_VERSION 1
//#define JLAB_ENERGY 1

struct rootDijetEvent {
    double W;
    double Q;
    double Pt;
    double qt;
    double z;
    double dphi;
    double deta;
    double Phi;
    double PtJet;
    double qtJet;
    double zJet;
    double PhiJet;
    double dphiJet;
    double detaJet;
    int nJets;
    int pol;
#if defined(NEW_VERSION)
    int n1;
    int n2;
#endif
};

double modulation(double *x, double*par) {
    double phi = x[0];
    double result = par[0]*(1+2*par[1]*cos(2*phi));
    return result;
}

void getv2(bool scaleErrors = true, bool usePolCut = false, int polCut = 0)
{
    //
    //  Files to analyze
    //
    vector<string> filenames;
    //filenames.push_back("5Mevt.root");
    //filenames.push_back("results.root");
    //filenames.push_back("eekt1M.root");
#if defined(JLAB_ENERGY)
    cout << "Using JLEIC file" << endl;
    filenames.push_back("results_lowE.root");
#else
    cout << "Using eRHIC file" << endl;
    filenames.push_back("eekt3M_20x100.root");
#endif
    //
    //  Create chain
    //
    TFile *f = 0;
    TTree *tree = 0;
    TChain *chain = 0;
    
    cout << "Setup new chain" << endl;
    chain = new TChain("tree"); // name of the tree is the argument
    for (auto name : filenames) {
        chain->Add(name.c_str());
        cout << "\tadded file '" << name << "' to chain." << endl;
    }
    
    tree = chain;
    
    double nEntries = tree->GetEntries();
    
    cout << "Available Entries = " << nEntries << endl;
    
    rootDijetEvent myEvent;
    TLorentzVector *mParton1 = 0;
    TLorentzVector *mParton2 = 0;
    TLorentzVector *mJet1 = 0;
    TLorentzVector *mJet2 = 0;
    
    tree->SetBranchAddress("event",&myEvent);
    tree->SetBranchAddress("mParton1",&mParton1);
    tree->SetBranchAddress("mParton2",&mParton2);
    tree->SetBranchAddress("mJet1",&mJet1);
    tree->SetBranchAddress("mJet2",&mJet2);
    
    tree->GetBranch("event");
    tree->GetBranch("mParton1");
    tree->GetBranch("mParton2");
    tree->GetBranch("mJet1");
    tree->GetBranch("mJet2");
    
    //
    //  Book histos
    //
    const int nbins = 31;
    TH1::SetDefaultSumw2(true);
    double const pi = acos(-1);
    TH1D *h_dNdphi_partons = new TH1D("h_dNdphi_partons", "partons: dN/dphi", nbins, 0, 2*pi);
    TH1D *h_dNdphiWeighted_partons = new TH1D("h_dNdphiWeighted_partons", "partons: dN/dphi weighted with cos(2phi)", nbins, 0, 2*pi);
    TH1D *h_dNdphi_jets = new TH1D("h_dNdphi_jets", "jets: dN/dphi", nbins, 0, 2*pi);
    TH1D *h_dNdphiWeighted_jets = new TH1D("h_dNdphiWeighted_jets", "jets: dN/dphi weighted with cos(2phi)", nbins, 0, 2*pi);
    h_dNdphi_partons->SetMinimum(0.);
    h_dNdphiWeighted_partons->SetMinimum(0.);
    h_dNdphi_jets->SetMinimum(0.);
    h_dNdphiWeighted_jets->SetMinimum(0.);
    
    //
    //  Loop events
    //
 
    for(int ievent=0; ievent< nEntries; ievent++) {
        
        if (ievent%100000 == 0) cout << "Processing event " << ievent << endl;
        tree->GetEntry(ievent);
        TLorentzVector mParton1Vec(*mParton1);
        TLorentzVector mParton2Vec(*mParton2);
        TLorentzVector mJet1Vec(*mJet1);
        TLorentzVector mJet2Vec(*mJet2);
        
        //
        //  ==> At this point all information of the tuple is available <==
        //
        
        //
        //  ==> Pt and qt corrections
        //
        myEvent.PtJet += 0.58;  // simple corrections instead of complicated unfolding
        
        const double lowPt = 2.75;
        const double highPt = 3.25;
        const double lowQt = 1.25;
        const double highQt = 1.75;
        if (myEvent.qt < highQt && myEvent.qt > lowQt && myEvent.Pt > lowPt && myEvent.Pt < highPt &&
            mParton1->PseudoRapidity() > 0 && mParton1->PseudoRapidity() < 3 &&
            mParton2->PseudoRapidity() > 0 && mParton2->PseudoRapidity() < 3) {
            if (!usePolCut || (usePolCut && myEvent.pol == polCut) ) {
                h_dNdphi_partons->Fill(myEvent.Phi,1.);
                h_dNdphiWeighted_partons->Fill(myEvent.Phi,cos(2*myEvent.Phi));
            }
        }
        if (myEvent.qtJet < highQt && myEvent.qtJet > lowQt && myEvent.PtJet > lowPt && myEvent.PtJet < highPt &&
            mJet1->PseudoRapidity() > 0 && mJet1->PseudoRapidity() < 3 &&
            mJet2->PseudoRapidity() > 0 && mJet2->PseudoRapidity() < 3) {
            if (!usePolCut || (usePolCut && myEvent.pol == polCut) ) {
                h_dNdphi_jets->Fill(myEvent.PhiJet,1.);
                h_dNdphiWeighted_jets->Fill(myEvent.PhiJet,cos(2*myEvent.PhiJet));
            }
        }
    } // end event loop
    
    //
    //  Calculate v2
    //
    double v2_partons = h_dNdphiWeighted_partons->Integral()/h_dNdphi_partons->Integral();
    double v2_jets = h_dNdphiWeighted_jets->Integral()/h_dNdphi_jets->Integral();
    
    cout << "partons v2: " << v2_partons << endl;
    cout << "jets v2: " << v2_jets << endl;
    
    //
    //  Rough estimate
    //  max/min = (1+2*v2)/(1-2*v2)
    //
    double ratio = h_dNdphi_partons->GetMaximum()/h_dNdphi_partons->GetMinimum();
    double v2Est = (ratio-1)/(2+2*ratio);
    cout << "partons v2 estimate: " << v2Est << endl;
    
    //
    //  Absolute normalization
    //
    const double nevents = nEntries;
    const double binwidth = h_dNdphi_partons->GetBinWidth(2);
    // for 15x100
    //double xsect_T = 6.00237e-06;  // was 1.63061e-05; // mb    pol = 0
    //double xsect_L = 3.07942e-06; //  was 7.54951e-06; // mb    pol = 1
#if defined(JLAB_ENERGY)
    cout << "Using JLEIC cross-sections" << endl;
    double xsect_T = 1.92114e-07; // mb     // 10x40
    double xsect_L = 9.22645e-08; // mb
#else
    cout << "Using eRHIC cross-sections" << endl;
    double xsect_T = 1.06558e-05; // mb     // 20x100
    double xsect_L = 5.18580e-06; // mb
#endif
    double scale = 1;
    
    if (usePolCut) {
        if (polCut == 1)
            scale = (xsect_L)/nevents/binwidth;
        else if (polCut==0)
            scale = (xsect_T)/nevents/binwidth;
    }
    else
        scale = (xsect_T+xsect_L)/nevents/binwidth;
        
    cout << "scale is " << scale << endl;
    TH1D *h_dNdphi_partons_norm = new TH1D((*h_dNdphi_partons)*scale);
    TH1D *h_dNdphi_jets_norm = new TH1D((*h_dNdphi_jets)*scale);
    h_dNdphi_partons_norm->SetMinimum(0.);
    h_dNdphi_jets_norm->SetMinimum(0.);
    
    //
    //   Errors corresponding to given integrated lumi
    //
    const double integratedLumi = 1; // 1/fb
    double cs_tot = 0;
    if (usePolCut) {
        if (polCut == 1)
            cs_tot = xsect_L;
        else if (polCut==0)
            cs_tot = xsect_T;
    }
    else
        cs_tot = xsect_T+xsect_L;
    // cs tot in mb to fb;
    cs_tot *= 1e12;
    double nTotalEvents = cs_tot * integratedLumi;  // expected # of events for given int. lumi
    double errorScale = sqrt(nTotalEvents/nEntries);  // factor to divide current errors by
    cout << " -> " << cs_tot << " fb;    N = " << cs_tot * nTotalEvents << endl;
    cout << nTotalEvents/nEntries << " more events than generated" << endl;
    cout << errorScale << " factor to scale error on entries (N)" << endl;
    
    if (scaleErrors) {
        for (int i=1; i<=nbins; i++) {
            double e = h_dNdphi_partons_norm->GetBinError(i);
            e /= errorScale;
            h_dNdphi_partons_norm->SetBinError(i, e);
            e = h_dNdphi_jets_norm->GetBinError(i);
            e /= errorScale;
            h_dNdphi_jets_norm->SetBinError(i, e);
        }
    }
    else
        cout << "No error scaling done due to given options." << endl;
    
    //
    //  General ROOT options
    //
    gROOT->SetStyle("Plain");
    gStyle->SetOptStat(11111);
    
    //
    //  1st Canvas
    //
    TCanvas *c1 = new TCanvas("c1"," ", 800, 800);
    c1->SetBorderSize(1);
    c1->SetFillColor(0);
    
    c1->Divide(2,2);
    
    char str[128];
    
    c1->cd(1);
    h_dNdphi_partons->Draw("E");
    sprintf(str,"v2 = %f", v2_partons);
    TText *t = new TText(0.15,0.8,str);
    t->SetNDC();
    t->SetTextColor(kRed);
    t->SetTextFont(43);
    t->SetTextSize(20);
    t->Draw();
    c1->cd(2);
    h_dNdphi_jets->Draw("E");
    sprintf(str,"v2 = %f", v2_jets);
    TText *t2 = new TText(0.15,0.8,str);
    t2->SetNDC();
    t2->SetTextColor(kRed);
    t2->SetTextFont(43);
    t2->SetTextSize(20);
    t2->Draw();
    c1->cd(3);
    h_dNdphi_partons_norm->SetMarkerStyle(20);
    h_dNdphi_partons_norm->SetMarkerColor(4);
    h_dNdphi_partons_norm->Draw("PE");
    h_dNdphi_partons_norm->GetXaxis()->SetTitle("#phi (rad)");
    h_dNdphi_partons_norm->GetYaxis()->SetTitle("d#sigma/d#phi (mb)");
    c1->cd(4);
    h_dNdphi_jets_norm->SetMarkerStyle(20);
    h_dNdphi_jets_norm->SetMarkerColor(2);
    h_dNdphi_jets_norm->Draw("PE");
    h_dNdphi_jets_norm->GetXaxis()->SetTitle("#phi (rad)");
    h_dNdphi_jets_norm->GetYaxis()->SetTitle("d#sigma/d#phi (mb)");
    //
    //   Fit sinus
    //
    TF1 *func2 = new TF1("func2", modulation, 0., 2*M_PI, 2);
    func2->SetParameters(1.e-9, 0.1);
    h_dNdphi_partons_norm->Fit(func2, "E");
    TF1 *func = new TF1("func", modulation, 0., 2*M_PI, 2);
    func->SetParameters(1.e-9, 0.1);
    h_dNdphi_jets_norm->Fit(func, "E");
    
    // afterthought
    c1->cd(3);
    h_dNdphi_jets_norm->Draw("PE same");
    
    string basename;
    if (usePolCut) {
        if (polCut == 0) basename = string("pol-0");
        if (polCut == 1) basename = string("pol-1");
    }
    else
        basename = string("pol-none");
    
    string epsfile = basename + string(".eps");
    
    cout << "Storing plots in '" << epsfile.c_str() << "'." << endl;
    c1->SaveAs(epsfile.c_str());
    
    string rfile = basename + string(".root");
    
    cout << "Storing histos in '" << rfile.c_str() << "'." << endl;
    TFile nfile(rfile.c_str(),"RECREATE");
    h_dNdphi_partons_norm->Write();
    h_dNdphi_jets_norm->Write();
    func->Write();
    func2->Write();
    nfile.Close();
    
    cout << "bye" << endl;
}