#include <TFile.h>
#include <TH2.h>
#include <TTree.h>
#include <iostream>

Float_t get_weight(TH2* histPR, TH2* histFR, Float_t elept, Float_t eleeta, Int_t eleistight) {
    // Find the bin indices corresponding to the given pT and eta values
    Float_t bin_pt_PR = histPR->GetXaxis()->FindBin(elept);
    Float_t bin_eta_PR = histPR->GetYaxis()->FindBin(fabs(eleeta));
    Float_t bin_pt_FR = histFR->GetXaxis()->FindBin(elept);
    Float_t bin_eta_FR = histFR->GetYaxis()->FindBin(fabs(eleeta));

    //std::cout << " pT " << elept <<" eta "<< eleeta << std::endl;
    //std::cout << " bins PR " << bin_pt_PR<<" and "<< bin_eta_PR << std::endl;
    //std::cout << " bins FR " << bin_pt_FR<<" and "<< bin_eta_FR << std::endl;
    
    // Get the content of the bin corresponding to the given pT and eta values
    Float_t value_p = histPR->GetBinContent(bin_pt_PR, bin_eta_PR);
    Float_t value_f = histFR->GetBinContent(bin_pt_FR, bin_eta_FR);

    //std::cout << value_p <<" and "<< value_f << std::endl;

    // Apply the 1-lepton formula to calculate weight
    Float_t weight = 1;
    if (eleistight > 0)
        weight = (value_p - 1) / (value_p - value_f);
    else
        weight = value_p / (value_p - value_f);
    return weight * value_f;
}

void apply_weight_to_data(const char* input_file_name, const char* output_file_name, TH2* histPR, TH2* histFR) {
    // Open input ROOT file
    TFile* input_file = TFile::Open(input_file_name);

    // Create output ROOT file
    TFile* output_file = new TFile(output_file_name, "RECREATE");

    // Get the tree containing the event data
    TTree* tree = (TTree*)input_file->Get("LJAlgo/nominal");


    // Create a new tree for the output with the weight branch added
    TTree* output_tree = tree->CloneTree(0);

    // Define variables to hold values from the tree
    Float_t elept;
    Float_t eleeta;
    Int_t eleistight;
    Bool_t lepton1_isSig;
    Bool_t lepton1_isEl;

    // Set up branches to read from the input tree
    tree->SetBranchAddress("lepton1_isSig", &eleistight);
    tree->SetBranchAddress("lepton1_pt", &elept);
    tree->SetBranchAddress("lepton1_eta", &eleeta);
    tree->SetBranchAddress("lepton1_isEl",&lepton1_isEl);

    // Add a new branch called "fakeweight" to the output tree
    Float_t fakeweight;
    output_tree->Branch("fakeweight", &fakeweight);
    output_tree->Branch("lepton1_isSig", &lepton1_isSig);

    // Loop over all events in the input tree
    for (int i = 0; i < tree->GetEntries(); ++i) {
        tree->GetEntry(i);
	if (lepton1_isEl <1 ) continue;
        // Calculate weight using the histograms
        fakeweight = get_weight(histPR, histFR, elept, eleeta, eleistight);
        if (i % 10000 == 0) {
	  std::cout << fakeweight << "tight?" << eleistight << std::endl;
          if (i>100000){
		break;
		}
        }
        lepton1_isSig = 1;
        // Fill the output tree with the weight added as a new branch
        output_tree->Fill();
    }

    // Write the output tree to the output file
    output_tree->Write();

    // Close input and output files
    input_file->Close();
    output_file->Close();
}

int Sample_creation_light() {
    // Open first ROOT file
    TFile* filePR = TFile::Open("PR_2016/output_PR.root");
    // Get the desired histogram from the first file
    TH2* histPR = (TH2*)filePR->Get("PR/PR_pT_eta_ele");

    // Open second ROOT file
    TFile* fileFR = TFile::Open("plots_2016/EleFR_jet360.root");
    // Get the desired histogram from the second file
    TH2* histFR = (TH2*)fileFR->Get("FR_pT_eta_EWKcorr_360");

    const char* inputfile = "/storage/agrp/roybr/leptonplusjet_skimmedtrees/data_13TeV/Merged/Run2_data_merged.root";

    // Apply weight to data in another ROOT file
    //apply_weight_to_data(inputfile, "Run2_fakes_datadriven_light.root", histPR, histFR);
    apply_weight_to_data(inputfile, "/storage/agrp/roybr/leptonplusjet_skimmedtrees/fakes/Run2_fakes_datadriven_light.root", histPR, histFR);

    return 0;
}