#include <iostream>
#include <TCanvas.h>
#include <vector>
//#include "../Settings.h"
//#include "fastjet/ClusterSequence.hh"
//#include "fastjet/contrib/SoftDrop.hh"
//#include "../LundGen.hh"
#include <cmath>
#include <TCanvas.h>
#include <TH2.h>
#include <TLorentzVector.h>
#include <TMath.h>
#include <TStyle.h>
#include <vector>
#include "/Users/josearias18/Desktop/QCDc2/pythia/Scripts/Reference/Header.h"
#include "RooAbsAnaConvPdf.h"
#include "RooAddPdf.h"
#include "RooArgusBG.h"
#include "RooCategory.h"
#include "RooConstVar.h"
#include "RooDataHist.h"
#include "RooDataSet.h"
#include "RooFFTConvPdf.h"
#include "RooFitResult.h"
#include "RooGaussian.h"
#include "RooGlobalFunc.h"
#include "RooPlot.h"
#include "RooPolynomial.h"
#include "RooRealSumFunc.h"
#include "RooRealVar.h"
#include "RooSimultaneous.h"
#include "TAxis.h"
#include "TFile.h"
#include "TH1.h"

using namespace std;
using namespace RooFit;

void Dzeroes_MassFit(int Ptbin = 3, bool Chi2 = true) {

    TString extension_read, extension_RootFiles;

    extension_RootFiles = TString("/Users/josearias18/Desktop/QCDc2/pythia/Root_Files/Djets/");

    extension_read = TString("Dzeroes_Variable_Tree_3067692");

    TFile fread(extension_RootFiles + extension_read + ".root", "READ");

    TString extension_out, pt_extension, extension_ptworkspace, extension_chi2;
    
    extension_chi2 = TString("chi2_0_"); 
    
    extension_out = TString("Dzeroes_MassFit");

    if (Ptbin == 0) {
        pt_extension = TString("_m");
        extension_ptworkspace = TString("_m");
    }

    if (Ptbin == 1) {
        pt_extension = TString("_m_20_30");
        extension_ptworkspace = TString("_m_20_30");
    }

    if (Ptbin == 2) {
        pt_extension = TString("_m_30_50");
        extension_ptworkspace = TString("_m_30_50");
    }

    if (Ptbin == 3) {
        pt_extension = TString("_m_50_100");
        extension_ptworkspace = TString("_m_50_100");
    }

    TString extension_tree;
    extension_tree = TString("Djet");

    TTree *Djet_tree;
    Djet_tree = (TTree *)fread.Get(extension_tree);
    int NumEvts = Djet_tree->GetEntries();

    cout << "Data entries: " << Djet_tree->GetEntries() << endl;

    TFile *f;
    
    if (Chi2) {
        f = new TFile(extension_RootFiles + extension_out + extension_chi2 + pt_extension + ".root", "RECREATE");
    } else {
        f = new TFile(extension_RootFiles + extension_out + pt_extension + ".root", "RECREATE");
    }

    double jet_pt, jet_eta;
    double jet_rap;
    double jet_px, jet_py, jet_pz, jet_e;
    double pi_px, pi_py, pi_pz, pi_e;
    double K_px, K_py, K_pz, K_e;

    double D0_ipchi2;
    int nSV;

    // TLorentzVector

    Djet_tree->SetBranchAddress("jet_pt", &jet_pt);
    Djet_tree->SetBranchAddress("jet_eta", &jet_eta);
    Djet_tree->SetBranchAddress("jet_rap", &jet_rap);

    Djet_tree->SetBranchAddress("jet_px", &jet_px);
    Djet_tree->SetBranchAddress("jet_py", &jet_py);
    Djet_tree->SetBranchAddress("jet_pz", &jet_pz);
    Djet_tree->SetBranchAddress("jet_e", &jet_e);

    Djet_tree->SetBranchAddress("D0_ipchi2", &D0_ipchi2);
    Djet_tree->SetBranchAddress("nSV", &nSV);

    Djet_tree->SetBranchAddress("pi_px", &pi_px);
    Djet_tree->SetBranchAddress("pi_py", &pi_py);
    Djet_tree->SetBranchAddress("pi_pz", &pi_pz);
    Djet_tree->SetBranchAddress("pi_e", &pi_e);

    Djet_tree->SetBranchAddress("K_px", &K_px);
    Djet_tree->SetBranchAddress("K_py", &K_py);
    Djet_tree->SetBranchAddress("K_pz", &K_pz);
    Djet_tree->SetBranchAddress("K_e", &K_e);


    int fitn = 1; ///DEFAULT - NO PRIOR FITTING (FASTEST)

//    int nBins = 40; ///NUMBER OF BINS - UP TO YOU.  NOTE: MORE BINS = MORE DETAILED DISTRIBUTION BUT LARGER UNCERTAINTIES; LESS BINS = BETTER UNCERTAINTIES BUT WORSE RESOLUTION
    Double_t rho = 1.; ///WIDTH OF GAUSSIANS USED IN REGRESSION FIT OF COMBINATORIC BACKGROUND
    
    double M_PDG = 1.86484;
    double M_width = 0.0075;
    double M_range = 8.*M_width;
    
    double LM = 1.805;
    double HM = 1.925; //1.92
    double MM, M_sigma, WM;
    double DM;
    int nBins;
    int nBins_bkg;

    ///// From Mass Fits (Literature)
    if (Ptbin == 0) {
        nBins = 120; //140
        MM = 1.86487;
        M_sigma = 0.0111;
        WM = 2. * M_sigma; //7 , 5

        DM = 0.015;
    }

    if (Ptbin == 1) {
        nBins = 120; //140
        MM = 1.86486;
        M_sigma = 0.0113;
        WM = 2. * M_sigma; //7

        DM = 0.015;
    }

    if (Ptbin == 2) {
        nBins = 80; //100
        MM = 1.86483;
        M_sigma = 0.0072;
        WM = 3. * M_sigma; //3

        DM = 0.015;
    }

    if (Ptbin == 3) {
        nBins = 40; //40
        nBins_bkg = 30;
        MM = 1.8649;
        M_sigma = 0.0073;
        WM = 2.5 * M_sigma; //2.5

        DM = 0.015;
    }


    TH1D *hchi2 = new TH1D("hchi2", "hchi2", nBins, -3., 5.);
    
    TH1D * hchi2_bkg; 
    if (Ptbin == 3) {
        hchi2_bkg = new TH1D("hchi2_bkg", "hchi2", nBins_bkg, -3., 5.);
    } else {
        hchi2_bkg = new TH1D("hchi2_bkg", "hchi2", nBins, -3., 5.);
    }

    TH1D *h1_mass = new TH1D("h1_mass", "h1_mass", nBins, LM, HM);

    int eventNum;
    double ipchi2_max;
    
    if (Chi2) {
        ipchi2_max = 0.0;
    } else {
        ipchi2_max = 5.0;
    }
    
    cout << "Requested # of events" << NumEvts << endl;

    for (int ev = 0; ev < NumEvts; ev++) {
        Djet_tree->GetEntry(ev);

        if (ev % 100000 == 0) {
            cout << "Executing event " << ev << endl;
        }

        TLorentzVector HFmeson, HFjet, pi, Kmeson, Jpsi;

        HFjet.SetPxPyPzE(jet_px, jet_py, jet_pz, jet_e);
        pi.SetPxPyPzE(pi_px, pi_py, pi_pz, pi_e);
        Kmeson.SetPxPyPzE(K_px, K_py, K_pz, K_e);

        HFmeson = pi + Kmeson;

        double dphi_HF_jet = checkphi(checkphi(HFmeson.Phi()) - checkphi(HFjet.Phi()));
        double dy_HF_jet = HFjet.Eta() - HFmeson.Rapidity();

        if (Ptbin == 0) {
            if (jet_pt < 20.0) {
                continue;
            }
        }

        if (Ptbin == 1) {
            if (jet_pt < 20.0) {
                continue;
            }

            if (jet_pt > 30.0) {
                continue;
            }
        }

        if (Ptbin == 2) {
            if (jet_pt < 30.0) {
                continue;
            }

            if (jet_pt > 50.0) {
                continue;
            }
        }

        if (Ptbin == 3) {
            if (jet_pt < 50.0) {
                continue;
            }

            if (jet_pt > 100.0) {
                continue;
            }
        }

//        if (log10(D0_ipchi2)>3.0) continue;

//        if (HFmeson.M() < LM || HFmeson.M() > HM) {
        if (HFmeson.M() < 1.80468 || HFmeson.M() > 1.925 ) {
            continue;
        }

        bool signal_cond = (HFmeson.M() > (M_PDG - WM) && (HFmeson.M() < (M_PDG + WM)) );
        bool bkg_cond = ( HFmeson.M() < 1.83 ) ||  ( HFmeson.M() > 1.895);

        // cout<<log10(D0_ipchi2);

        if (bkg_cond) {
            hchi2_bkg->Fill(log10(D0_ipchi2));
        }

        if (signal_cond) {
            hchi2->Fill(log10(D0_ipchi2));
        }
        
        hchi2->Fill(log10(D0_ipchi2));

        if (log10(D0_ipchi2) > ipchi2_max ) {
                continue;
            }
        
        h1_mass->Fill(HFmeson.M());
    }

//  TLatex latex;
    ///fitn SELECTS STRATEGY FOR FITTING OF COMBINATORIC BACKGROUND///
    //int fitn = 1; ///FIT TO BACKGROUND WITH INITIAL HESSE TURNED ON
    //int fitn = 0; ///STANDARD FIT TO BACKGROUND

    ///CHANGE FOR D0 MASS SIDEBAND
    TCanvas *cm = new TCanvas("cm", "c", 700, 1000);
    cm->Divide(1, 2);
    cm->cd(1);
    ////////Chi2 Distributions////////
    RooRealVar HFMass("HFMass", "HFMass", LM, HM); ///LOG(CHI_IP^2) VARIABLE THAT IS FED INTO ROOFIT FUNCTIONS
    HFMass.setBins(nBins);
    RooDataHist D0Mass("D0Mass", "D0Mass", RooArgList(HFMass), h1_mass, 1.);

     RooRealVar Mean("Mean","",M_PDG, M_PDG-DM, M_PDG+DM); //
//    RooRealVar Mean("Mean", "", 1.86, 1.87);
    //RooRealVar CBS("CBS","CBS",DM/2.,4./1000.,DM);
    RooRealVar CBS("CBS", "CBS", 4. / 1000., DM);
    RooRealVar n("n", "n", 1.);
    //RooRealVar a("a","a",2.,1.,5.);
    RooRealVar a("a", "a", 1., 5.);
    RooCBShape sig1("cb", "Signal", HFMass, Mean, CBS, a, n);

    //Sigma,f,a are obtained from fit for kinematically unbinned data, and I do the fit outside the file.
    // RooFormulaVar Sigma("Sigma","1.647*CBS",RooArgSet(CBS));
    //RooRealVar Sigma("Sigma", "Sigma", DM/2.,6./1000., DM);
    RooRealVar Sigma("Sigma", "Sigma", 6. / 1000., DM);
    RooGaussian gaus("gaus", "", HFMass, Mean, Sigma);
    RooRealVar frac;
    
    if ( Ptbin == 0 )
    {
        frac = RooRealVar("frac", "", 0.3, 0.999); //0.1
        
    } //0.1, 0.34 (30-50)
    if ( Ptbin == 1 )
    {
        frac = RooRealVar("frac", "", 0.1, 0.999); //0.1
        
    }
    if ( Ptbin == 2 )
    {
        frac = RooRealVar("frac", "", 0.34, 0.999);
        
    }
    if ( Ptbin == 3 )
    {
        frac = RooRealVar("frac", "", 0.3, 0.999);
        
    }

    // f.setVal(0.667);f.setConstant();
    // a.setVal(2.444);a.setConstant();
    RooAddPdf sig("sig", "gau1+gau2", RooArgSet(sig1, gaus), RooArgSet(frac));

    RooRealVar ns("ns", "nSignal", 0., 3067692.); //5000.,
    RooExtendPdf sign("sign", "sign", sig, ns);

    //BKG component
    RooRealVar p("p","p",-0.,-1./(M_PDG+M_range),1./(M_PDG+M_range));
//    RooRealVar p("p", "p", -1., 1.);
	RooPolynomial bkg("bkg","bkg",HFMass,RooArgSet(p));
//    RooChebychev bkg("bkg", "bkg", HFMass, RooArgSet(p));
    RooRealVar nb("nb", "nBKG", 0., 3067692.);  //300.,
    RooExtendPdf bkgn("bkgn", "bkgn", bkg, nb);

    //sig+bkg

    RooAddPdf model("model", "sig+bkg", RooArgSet(bkgn, sign));

    RooArgSet *params = model.getParameters(HFMass);
    // Save snapshot of prefit parameters
    RooArgSet *initParams = (RooArgSet *)params->snapshot();

    model.fitTo(D0Mass);

    /////////////////Plotting/////////////////
    double binsize = (HM - LM) / (double)nBins;
    RooPlot *xframe = HFMass.frame(Title(Form(";M_{K#pi} [GeV];Events/(%.1f MeV) ", binsize * 1000.)));
    D0Mass.plotOn(xframe, Name("D0Mass"), MarkerSize(0.81));

    model.plotOn(xframe, Name("bkgn"), Components(bkgn), LineColor(kBlue), LineStyle(kDashed), PrintEvalErrors(-1));
    model.plotOn(xframe, Name("sign"), Components(sign), LineColor(kGreen), LineStyle(5), PrintEvalErrors(-1));
    model.plotOn(xframe, Name("model"), LineWidth(2), LineColor(kRed), PrintEvalErrors(-1));
    double chi2ndf = xframe->chiSquare(11);
    // model.plotOn(xframe, Components(sec), FillColor(kRed), FillStyle(3002), DrawOption("F"),PrintEvalErrors(-1));
    // model.plotOn(xframe, Components(comb), FillColor(kGreen+2), FillStyle(3003), DrawOption("F"),PrintEvalErrors(-1));

    xframe->Draw("SAME");
    TLegend *leg1 = new TLegend(0.61,0.55,0.88,0.87);
    // leg1->SetFillColor(kWhite);
    // leg1->SetLineColor(kWhite);
    leg1->SetBorderSize(0);
    if (Ptbin == 0) {
        leg1->AddEntry("D0Mass", "Data p_{T} > 20 GeV/c", "PE");
    }

    if (Ptbin == 1) {
        leg1->AddEntry("D0Mass", "Data 20 GeV/c < p_{T} < 30 GeV/c", "PE");
    }

    if (Ptbin == 2) {
        leg1->AddEntry("D0Mass", "Data 30 GeV/c < p_{T} < 50 GeV/c", "PE");
    }

    if (Ptbin == 3) {
        leg1->AddEntry("D0Mass", "Data 50 GeV/c < p_{T} < 100 GeV/c", "PE");
    }
    
    leg1->AddEntry("model", "Fit model", "L");
    leg1->AddEntry("sign", "Signal", "L");
    leg1->AddEntry("bkgn", "Comb. bkg.", "L");
    leg1->AddEntry((TObject *)0, Form("#chi^{2}/ndf = %.2f", chi2ndf), "");
    leg1->AddEntry((TObject *)0, Form(" #chi^{2}_{IP, D^{0}} < %.2f", ipchi2_max), "");
    leg1->AddEntry((TObject *)0, Form("m_{ D^{0}} = %.2f MeV/c^{2}", Mean.getVal()*1000.), "");
//    leg1->AddEntry((TObject *)0, Form("n_{s} = %.0f ", ns.getVal()), "");
//    leg1->AddEntry((TObject *)0, Form("n_{b} = %.0f ", nb.getVal()), "");

    leg1->Draw("SAME");
    // res->Print();

    if (Chi2) {
        cm->SaveAs("/Users/josearias18/Desktop/QCDc2/pythia/Plots/Djets/D0_MassFit" + extension_chi2 + pt_extension + ".pdf");
    } else {
        cm->SaveAs("/Users/josearias18/Desktop/QCDc2/pythia/Plots/Djets/D0_MassFit" + pt_extension + ".pdf");
    }
    
    cout << "Chi2/ndf = " << chi2ndf << endl;

//    HFMass.setRange("signal", MM - 0.02, MM + 0.02);
    HFMass.setRange("signal", MM - WM, MM + WM);
    RooAbsReal *intsigX = sign.createIntegral(HFMass, NormSet(HFMass), Range("signal"));

    RooWorkspace *w = new RooWorkspace("w", "workspace");
    // w->import(sign);
    // w->import(bkgn);
    w->import(model);
    w->import(D0Mass);
    w->Print();
    
    if (Chi2) {
        w->writeToFile(  extension_RootFiles + "workspace_Dmassfit" +  extension_chi2 + extension_ptworkspace + ".root");
    } else {
        w->writeToFile(  extension_RootFiles + "workspace_Dmassfit" + extension_ptworkspace + ".root");
    }
//    w->writeToFile(  extension_RootFiles + "workspace_Dmassfit" + extension_ptworkspace + ".root");
    if (Chi2) {
        params->printLatex(Sibling(*initParams), OutputFile("/Users/josearias18/Desktop/QCDc2/pythia/RooFit_tex_Files/D0_MassFit" +  extension_chi2 + pt_extension+ ".tex"));
    } else {
        params->printLatex(Sibling(*initParams), OutputFile("/Users/josearias18/Desktop/QCDc2/pythia/RooFit_tex_Files/D0_MassFit" + pt_extension+ ".tex"));
    }
//    params->printLatex(Sibling(*initParams), OutputFile("/Users/josearias18/Desktop/QCDc2/pythia/RooFit_tex_Files/D0_MassFit" + pt_extension+ ".tex"));
    
//    f.Write();
//    f.Close();
    f->Write();
    f->Close();
    
    delete f;
    delete w;
    
//    double rec[2] = {
//        ns.getValV(), ns.getError()
//    };
//    intsigX->Print();
//    return rec;
    
   
}