#include <iostream>
#include <fstream>
#include <iomanip>
#include <string>
#include <vector>
#include <sstream>
#include "TFile.h"
#include "TNtuple.h"
#include "TTree.h"
#include "TH2F.h"
#include "TH1F.h"
#include "TMath.h"
#include "TChain.h"
#include "TLorentzVector.h"
#include "TLegend.h"
#include "TVector3.h"
#include "TStyle.h"
#include "TF1.h"
#include "TCanvas.h"
#include "TLine.h"
#include <TRandom3.h>
#include "TGraphErrors.h"
#include "TLatex.h"
#include "TPad.h"
#include "TStyle.h"
#include <numeric>
#include <cmath>

void phiMass_t_pol_cut() {
    TFile *filein = new TFile("flat_pi0pippim__B4_cut.root");
    TTree *tree = (TTree *)filein->Get("pi0pippim__B4_cut");

    UInt_t run;
    Int_t pol;
    float tagWeight, Mandelstam_t;

    TLorentzVector *pim_p4_kin = 0;
    TLorentzVector *pip_p4_kin = 0;
    TLorentzVector *decaypi0_p4_kin = 0;

    tree->SetBranchAddress("decaypi0_p4_kin", &decaypi0_p4_kin);
    tree->SetBranchAddress("pim_p4_kin", &pim_p4_kin);
    tree->SetBranchAddress("pip_p4_kin", &pip_p4_kin);

    tree->SetBranchAddress("run", &run);
    tree->SetBranchAddress("tagWeight", &tagWeight);
    tree->SetBranchAddress("pol", &pol);
    tree->SetBranchAddress("Mandlestam_t", &Mandelstam_t);

    //define the histograms loop

    ostringstream histName, histTitle;

    int t_bins = 10;

    TH1F *h_PhiMass_pol0_t_[t_bins];
    TH1F *h_PhiMass_pol90_t_[t_bins];
    TH1F *h_PhiMass_pol45_t_[t_bins];
    TH1F *h_PhiMass_pol135_t_[t_bins];


    for (int j = 0; j < t_bins; j++) {
            
        histName.str("");
        histTitle.str("");

        float t_lower = j * 0.1;
        float t_upper = (j + 1) * 0.1;
        
        if (pol == 0) {
            histName << "h_PhiMass_pol0_t" << std::setw(2) << std::setfill('0') << j << std::setw(2) << std::setfill('0') << j + 1;
            histTitle << std::fixed << std::setprecision(2) << t_lower << " < -t < " << t_upper << ";M(#phi) [GeV/c^{2}]";
            h_PhiMass_pol0_t_[j] = new TH1F(histName.str().c_str(), histTitle.str().c_str(), 300, 0.8, 1.2);
        }

        if (pol == 90) {
            histName << "h_PhiMass_pol90_t" << std::setw(2) << std::setfill('0') << j << std::setw(2) << std::setfill('0') << j + 1;
            histTitle << std::fixed << std::setprecision(2) << t_lower << " < -t < " << t_upper << ";M(#phi) [GeV/c^{2}]";
            h_PhiMass_pol90_t_[j] = new TH1F(histName.str().c_str(), histTitle.str().c_str(), 300, 0.8, 1.2);
        }

        if (pol == 45) {
            histName << "h_PhiMass_pol45_t" << std::setw(2) << std::setfill('0') << j << std::setw(2) << std::setfill('0') << j + 1;
            histTitle << std::fixed << std::setprecision(2) << t_lower << " < -t < " << t_upper << ";M(#phi) [GeV/c^{2}]";
            h_PhiMass_pol45_t_[j] = new TH1F(histName.str().c_str(), histTitle.str().c_str(), 300, 0.8, 1.2);
        }

        if (pol == 135) {
            histName << "h_PhiMass_pol135_t" << std::setw(2) << std::setfill('0') << j << std::setw(2) << std::setfill('0') << j + 1;
            histTitle << std::fixed << std::setprecision(2) << t_lower << " < -t < " << t_upper << ";M(#phi) [GeV/c^{2}]";
            h_PhiMass_pol135_t_[j] = new TH1F(histName.str().c_str(), histTitle.str().c_str(), 300, 0.8, 1.2);
        } 
    }


    TStyle *dStyle = new TStyle("JLab_Style", "JLab_Style");
    dStyle->Reset("Modern");

    TLorentzVector Pi0P4, PiMinusP4, PiPlusP4, PhiP4;

    Long64_t NEntries = tree->GetEntries();
    //NEntries = 400000;

    cout << "There are " << NEntries << " entries." << endl;

    for (Long64_t i_entry = 0; i_entry < NEntries; i_entry++) {
        cout << i_entry << endl;
        
        tree->GetEntry(i_entry);

        if (i_entry % 400000 == 0) {
            cout << i_entry << " events processed" << endl;
        }

        PiMinusP4.SetPxPyPzE(pim_p4_kin->Px(), pim_p4_kin->Py(), pim_p4_kin->Pz(), pim_p4_kin->E());
        PiPlusP4.SetPxPyPzE(pip_p4_kin->Px(), pip_p4_kin->Py(), pip_p4_kin->Pz(), pip_p4_kin->E());
        Pi0P4.SetPxPyPzE(decaypi0_p4_kin->Px(), decaypi0_p4_kin->Py(), decaypi0_p4_kin->Pz(), decaypi0_p4_kin->E());

        PhiP4 = PiMinusP4 + PiPlusP4 + Pi0P4;

        double PhiMass = PhiP4.M();
        

        for (int j = 0; j < t_bins; j++) {
    
            float t_lower = j * 0.1;
            float t_upper = (j + 1) * 0.1;

            if (-Mandelstam_t > t_lower && -Mandelstam_t < t_upper) {

                cout << "" << endl;  
                cout << "event " << i_entry << " filled" << endl;
                cout << "Mass: " << PhiMass << endl;
                cout << "pol: " << pol << endl;
                cout << "-t: " << -Mandelstam_t << endl;

                if (pol == 0) {
                    h_PhiMass_pol0_t_[j]->Fill(PhiMass, tagWeight);
                } else if (pol == 90) {
                    h_PhiMass_pol90_t_[j]->Fill(PhiMass, tagWeight);
                } else if (pol == 45) {
                    h_PhiMass_pol45_t_[j]->Fill(PhiMass, tagWeight);
                } else if (pol == 135) {
                    h_PhiMass_pol135_t_[j]->Fill(PhiMass, tagWeight);
                }
            }     
        }    
    }

    // Size & Margins
    dStyle->SetCanvasDefH(800);
    dStyle->SetCanvasDefW(1200);

    // Basic object fill colors
    dStyle->SetCanvasBorderMode(0);
    dStyle->SetLegendFillColor(0);
    dStyle->SetCanvasColor(0);

    // Stat box
    dStyle->SetOptFit(112);
    dStyle->SetOptStat(000001111); // entries only //mean/rms should be replaced with a fit

    // Titles and Labels
    dStyle->SetLabelSize(0.07, "xyz");
    dStyle->SetTitleSize(0.075, "xyz");
    dStyle->SetTitleOffset(1.1, "x");
    dStyle->SetTitleOffset(0.8, "y");

    // Margins (correlated with title/label size)
    dStyle->SetPadBottomMargin(0.15);

    // Default Histogram Style Settings
    // dStyle->SetHistFillColor(kTeal);
    dStyle->SetHistFillStyle(1001); // solid fill (0 for hollow) //see http://root.cern.ch/root/html/TAttFill.html
    dStyle->SetHistLineColor(kBlack);
    dStyle->SetHistMinimumZero(kTRUE);

    // Default Function Style Settings
    dStyle->SetFuncColor(kBlack);
    dStyle->SetFuncStyle(1); // see http://root.cern.ch/root/html/TAttLine.html
    dStyle->SetFuncWidth(3);

    // Default Line Style Settings
    dStyle->SetLineColor(kBlack);
    dStyle->SetLineStyle(1); // see http://root.cern.ch/root/html/TAttLine.html
    dStyle->SetLineWidth(1); // is border of TLegend

    // Default Graph Style Settings
    dStyle->SetMarkerColor(kBlack);
    dStyle->SetMarkerSize(1);   // see http://root.cern.ch/root/html/TAttMarker.html
    dStyle->SetMarkerStyle(22); // see http://root.cern.ch/root/html/TAttMarker.html

    // Grid
    dStyle->SetPadGridX(kTRUE);
    dStyle->SetPadGridY(kTRUE);

    // Histogram/Graph Title Size
    dStyle->SetTitleSize(0.075, "t"); //"t": could be anything that is not "x" "y" or "z"

    // Palette
    dStyle->SetNumberContours(200);
    #if ROOT_VERSION_CODE >= ROOT_VERSION(6, 0, 0)
        dStyle->SetPalette(kBird); // 57
    #else
        dStyle->SetPalette(55); // rainbow
    #endif
        dStyle->cd();
        TLatex ytitle;
        TLatex maintitle;
        TLatex xtitle;
        xtitle.SetTextFont(42);
        ytitle.SetTextFont(42);
        gStyle->SetOptStat(0000);
    

    // draw dem histies
    TCanvas *c1 = new TCanvas("c1", "Phi Masses", 800, 600);
    c1->cd(1);
    h_PhiMass_pol0_t_[0]->Draw();
}