// BDT_Classification_FromHistograms.cpp
#include "TFile.h"
#include "TH1.h"
#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#include "TMVA/DataLoader.h"
#include <vector>
#include <iostream>
using namespace std;
int main() {
    // Initialize TMVA
    TMVA::Tools::Instance();
    auto outputFile = TFile::Open("TMVA_BDT_Results_FromHistograms.root", "RECREATE");

    // Create a factory and a DataLoader for handling data
    TMVA::Factory factory("TMVAClassification", outputFile, "!V:!Silent:Color:DrawProgressBar:AnalysisType=Classification");
    TMVA::DataLoader loader("dataset");

    // Add discriminating variables
    loader.AddVariable("b1_pt", 'F');
    loader.AddVariable("b1_eta", 'F');
    loader.AddVariable("b2_pt", 'F');
    loader.AddVariable("b2_eta", 'F');
    loader.AddVariable("h_mass", 'F');
    loader.AddVariable("h_pt", 'F');
    loader.AddVariable("h_eta", 'F');
    loader.AddVariable("DR_b", 'F');
    loader.AddVariable("jet_pt", 'F');
    loader.AddVariable("jet_eta", 'F');
    loader.AddVariable("jet_phi", 'F');
    loader.AddVariable("jet_ET", 'F');
    loader.AddVariable("jet_mass", 'F');
    loader.AddVariable("jet_DeltaEta", 'F');
    loader.AddVariable("jet_Deltaphi", 'F');
    loader.AddVariable("jet_EhadOverEem", 'F');
    loader.AddVariable("missingET_MET", 'F');
    loader.AddVariable("missingEt_Phi", 'F');
    loader.AddVariable("scalarHT_HT", 'F');
    loader.AddVariable("jet_bTag", 'I');
    
    // Load input file with histograms
    TFile* inputFile = TFile::Open( "/home/ahmeedsayed/TMVAHist/h90BP1_bb.root");
  // TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h90BP2_bb.root");
  // TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h95BP1_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h95BP3_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h96BP1_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h96BP2_bb.root");
  // TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h96BP3_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h96BP6_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/h_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/z_bb.root");
   //TFile* inputFile = TFile::Open("/home/ahmeedsayed/TMVAHist/bb.root");
  
   
    if (!inputFile || inputFile->IsZombie()) {
        std::cerr << "Error opening input file with histograms!" << std::endl;
        return 1;
    }

    // Retrieve histograms by name
     TH1D* hist_b1_pt = (TH1D*)inputFile->Get("b1_pt");
    TH1D* hist_b1_eta = (TH1D*)inputFile->Get("b1_eta");
    TH1D* hist_b2_pt = (TH1D*)inputFile->Get("b2_pt");
    TH1D* hist_b2_eta = (TH1D*)inputFile->Get("b2_eta");
    TH1D* hist_h_mass = (TH1D*)inputFile->Get("h_mass");
    TH1D* hist_h_pt = (TH1D*)inputFile->Get("h_pt");
     TH1D* hist_h_eta = (TH1D*)inputFile->Get("h_eta");
    TH1D* hist_DR_b = (TH1D*)inputFile->Get("DR_b");
    TH1D* hist_jet_pt = (TH1D*)inputFile->Get("jet_pt");
    TH1D* hist_jet_eta = (TH1D*)inputFile->Get("jet_eta");
    TH1D* hist_jet_phi = (TH1D*)inputFile->Get("jet_phi");
    TH1D* hist_jet_ET = (TH1D*)inputFile->Get("jet_ET");
    TH1D* hist_jet_mass = (TH1D*)inputFile->Get("jet_mass");
    TH1D* hist_jet_DeltaEta= (TH1D*)inputFile->Get("jet_DeltaEta");
     TH1D* hist_jet_Deltaphi = (TH1D*)inputFile->Get("jet_Deltaphi");
    TH1D* hist_jet_EhadOverEem = (TH1D*)inputFile->Get("jet_EhadOverEem");
    TH1D* hist_missingET_MET = (TH1D*)inputFile->Get("missingET_MET");
    TH1D* hist_missingEt_Phi = (TH1D*)inputFile->Get("missingEt_Phi");
    TH1D* hist_scalarHT_HT = (TH1D*)inputFile->Get("scalarHT_HT");
    TH1I* hist_jet_bTag = (TH1I*)inputFile->Get("jet_bTag");
    // Verify histograms exist
   
   
    if (!hist_b1_pt || !hist_b1_eta || !hist_b2_pt || !hist_b2_eta || !hist_h_mass || !hist_h_pt || !hist_h_eta ||
        !hist_DR_b || !hist_jet_pt || !hist_jet_eta || !hist_jet_phi || !hist_jet_ET || !hist_jet_mass ||
        !hist_jet_DeltaEta || !hist_jet_Deltaphi || !hist_jet_EhadOverEem || !hist_missingET_MET ||
        !hist_missingEt_Phi || !hist_scalarHT_HT || !hist_jet_bTag) {
        std::cerr << "One or more histograms not found in the file!" << std::endl;
   
    }

    // Convert histogram data into vectors for TMVA training
    
    int nEntries =hist_b1_pt->GetNbinsX(); // Assuming the same number of bins for simplicity
    for (int i = 1; i <= nEntries; ++i) {  // Start at 1 since bin 0 is underflow in ROOT histograms
       
        float b1_pt = hist_b1_pt->GetBinContent(i);
        float b1_eta = hist_b1_eta->GetBinContent(i);
        float b2_pt = hist_b2_pt->GetBinContent(i);
        float b2_eta = hist_b2_eta->GetBinContent(i);
        float h_mass =hist_h_mass->GetBinContent(i);
        float h_pt =hist_h_pt->GetBinContent(i);
        float h_eta = hist_h_eta->GetBinContent(i);
        float DR_b =hist_DR_b->GetBinContent(i);
        float jet_pt =  hist_jet_pt->GetBinContent(i);
        float jet_eta = hist_jet_eta->GetBinContent(i);
        float jet_phi = hist_jet_phi->GetBinContent(i);
        float jet_ET = hist_jet_ET->GetBinContent(i);
        float jet_mass = hist_jet_mass->GetBinContent(i);
        float jet_DeltaEta = hist_jet_DeltaEta->GetBinContent(i);
         float jet_Deltaphi = hist_jet_Deltaphi->GetBinContent(i);
         float jet_EhadOverEem = hist_jet_EhadOverEem->GetBinContent(i);
         float missingET_MET = hist_missingET_MET->GetBinContent(i);
        float missingEt_Phi = hist_missingEt_Phi->GetBinContent(i);
        float scalarHT_HT = hist_scalarHT_HT->GetBinContent(i);
        int jet_bTag = hist_jet_bTag->GetBinContent(i);
       
       
         // Add data to the DataLoader
       if(b1_pt && b1_eta && b2_pt && b2_eta && h_mass && h_pt && h_eta && DR_b && jet_pt && jet_eta && jet_phi &&jet_ET &&
                                jet_mass && jet_DeltaEta && jet_Deltaphi && jet_EhadOverEem && missingET_MET && missingEt_Phi && scalarHT_HT && jet_bTag > 0)
            
         {
          loader.AddSignalTrainingEvent({b1_pt, b1_eta, b2_pt, b2_eta,h_mass,h_pt,h_eta,DR_b,jet_pt,jet_eta,jet_phi,jet_ET,
                                jet_mass, jet_DeltaEta, jet_Deltaphi, jet_EhadOverEem, missingET_MET, missingEt_Phi, scalarHT_HT, (float)jet_bTag});
        
          loader.AddSignalTestEvent({b1_pt, b1_eta, b2_pt, b2_eta,h_mass,h_pt,h_eta,DR_b,jet_pt,jet_eta,jet_phi,jet_ET,
                                 jet_mass, jet_DeltaEta, jet_Deltaphi, jet_EhadOverEem, missingET_MET, missingEt_Phi, scalarHT_HT, (float)jet_bTag});
       
   }
        
          if(b1_pt && b1_eta && b2_pt && b2_eta && h_mass && h_pt && h_eta && DR_b && jet_pt && jet_eta && jet_phi &&jet_ET &&
                                jet_mass && jet_DeltaEta && jet_Deltaphi && jet_EhadOverEem && missingET_MET && missingEt_Phi && scalarHT_HT && jet_bTag > 0)
       
       {
       loader.AddBackgroundTrainingEvent({b1_pt, b1_eta, b2_pt, b2_eta,h_mass,h_pt,h_eta,DR_b,jet_pt,jet_eta,jet_phi,jet_ET,
                                jet_mass, jet_DeltaEta, jet_Deltaphi, jet_EhadOverEem, missingET_MET, missingEt_Phi, scalarHT_HT, (float)jet_bTag});
                     
       loader.AddBackgroundTestEvent({b1_pt, b1_eta, b2_pt, b2_eta,h_mass,h_pt,h_eta,DR_b,jet_pt,jet_eta,jet_phi,jet_ET,
                              jet_mass, jet_DeltaEta, jet_Deltaphi, jet_EhadOverEem, missingET_MET, missingEt_Phi, scalarHT_HT, (float)jet_bTag});
      
   }
    
     }
    
      TCut mycut = "abs(b1_pt > 40 && b2_pt > 30 && h_mass > 40 && h_pt > 40 && h_eta < 10 && scalarHt_Ht > 60)";
  loader.PrepareTrainingAndTestTree( mycut, 
                                         "nTrain_Signal=3000:nTrain_Background=3000:SplitMode=Random:NormMode=NumEvents:!V" );

    
    
    // Setup BDT configuration
    factory.BookMethod(&loader, TMVA::Types::kBDT, "BDT",
                       "!H:!V:NTrees=100:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:SeparationType=GiniIndex:nCuts=20");

    // Train the model
    factory.TrainAllMethods();
    factory.TestAllMethods();
    factory.EvaluateAllMethods();

    // Clean up
    outputFile->Close();
    inputFile->Close();

    std::cout << "BDT training completed using histograms!" << std::endl;
    return 0;
}

// Add this at the end of BDT_Classification.cpp

  int BDT_Classification() {
    return main();
}