#include "TFile.h"
#include "TH1.h"
#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#include "TMVA/DataLoader.h"
#include <vector>
#include <memory>
#include <iostream>

    void BDT_Classification() {
    // 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 with correct types
    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_bTag", 'I');
    loader.AddVariable("jet_EhadOverEem", 'F');
    loader.AddVariable("missingET_MET", 'F');
    loader.AddVariable("missingEt_Phi", 'F');
    loader.AddVariable("scalarHT_HT", 'F');

    // List of file paths
    std::vector<std::string> filePaths = {
        "/home/ahmeedsayed/TMVAHist/h90BP1_bb.root",
        "/home/ahmeedsayed/TMVAHist/h90BP2_bb.root",
        "/home/ahmeedsayed/TMVAHist/h95BP1_bb.root",
        "/home/ahmeedsayed/TMVAHist/h95BP3_bb.root",
        "/home/ahmeedsayed/TMVAHist/h96BP1_bb.root",
        "/home/ahmeedsayed/TMVAHist/h96BP2_bb.root",
        "/home/ahmeedsayed/TMVAHist/h96BP3_bb.root",
        "/home/ahmeedsayed/TMVAHist/h96BP6_bb.root",
        "/home/ahmeedsayed/TMVAHist/h_bb.root",
        "/home/ahmeedsayed/TMVAHist/z_bb.root",
        "/home/ahmeedsayed/TMVAHist/bb.root"
    };

    // Vector to store pointers to TFile objects
    std::vector<std::unique_ptr<TFile>> inputFiles;

    // Open each file and store it in the vector
    for (const auto& filePath : filePaths) {
        auto file = std::unique_ptr<TFile>(TFile::Open(filePath.c_str()));
        if (!file || file->IsZombie()) {
            std::cerr << "Error opening input file: " << filePath << std::endl;
            return ; // Exiting function on file open failure
        }
        inputFiles.push_back(std::move(file));
    }

    // Loop through each file and retrieve histograms by name
    for (const auto& file : inputFiles) {
        std::string fileName = file->GetName();

        // Retrieve histograms
        TH1F* hist_b1_pt = dynamic_cast<TH1F*>(file->Get("b1_pt"));
        TH1F* hist_b1_eta = dynamic_cast<TH1F*>(file->Get("b1_eta"));
        TH1F* hist_b2_pt = dynamic_cast<TH1F*>(file->Get("b2_pt"));
        TH1F* hist_b2_eta = dynamic_cast<TH1F*>(file->Get("b2_eta"));
        TH1F* hist_h_mass = dynamic_cast<TH1F*>(file->Get("h_mass"));
        TH1F* hist_h_pt = dynamic_cast<TH1F*>(file->Get("h_pt"));
        TH1F* hist_h_eta = dynamic_cast<TH1F*>(file->Get("h_eta"));
        TH1F* hist_DR_b = dynamic_cast<TH1F*>(file->Get("DR_b"));
        TH1F* hist_jet_pt = dynamic_cast<TH1F*>(file->Get("jet_pt"));
        TH1F* hist_jet_eta = dynamic_cast<TH1F*>(file->Get("jet_eta"));
        TH1F* hist_jet_phi = dynamic_cast<TH1F*>(file->Get("jet_phi"));
        TH1F* hist_jet_ET = dynamic_cast<TH1F*>(file->Get("jet_ET"));
        TH1F* hist_jet_mass = dynamic_cast<TH1F*>(file->Get("jet_mass"));
        TH1F* hist_jet_DeltaEta = dynamic_cast<TH1F*>(file->Get("jet_DeltaEta"));
        TH1F* hist_jet_Deltaphi = dynamic_cast<TH1F*>(file->Get("jet_Deltaphi"));
        TH1I* hist_jet_bTag = dynamic_cast<TH1I*>(file->Get("jet_bTag"));
        TH1F* hist_jet_EhadOverEem = dynamic_cast<TH1F*>(file->Get("jet_EhadOverEem"));
        TH1F* hist_missingET_MET = dynamic_cast<TH1F*>(file->Get("missingET_MET"));
        TH1F* hist_missingEt_Phi = dynamic_cast<TH1F*>(file->Get("missingEt_Phi"));
        TH1F* hist_scalarHT_HT = dynamic_cast<TH1F*>(file->Get("scalarHT_HT"));

        // Check if essential histograms were successfully retrieved
        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_bTag || !hist_jet_EhadOverEem || !hist_missingET_MET ||
            !hist_missingEt_Phi || !hist_scalarHT_HT) {
            std::cerr << "Error retrieving histogram(s) from file: " << fileName << std::endl;
            continue;
        }

        int nEntries = hist_b1_pt->GetNbinsX();
        if (nEntries == 0) {
            std::cerr << "Warning: No entries in histograms for file: " << fileName << std::endl;
            continue;
        }

        // Loop through each bin as an event
        for (int i = 1; i <= nEntries; ++i) {
            std::vector<Double_t> values = {
                hist_b1_pt->GetBinCenter(i), hist_b1_eta->GetBinCenter(i),
                hist_b2_pt->GetBinCenter(i), hist_b2_eta->GetBinCenter(i),
                hist_h_mass->GetBinCenter(i), hist_h_pt->GetBinCenter(i),
                hist_h_eta->GetBinCenter(i), hist_DR_b->GetBinCenter(i),
                hist_jet_pt->GetBinCenter(i), hist_jet_eta->GetBinCenter(i),
                hist_jet_phi->GetBinCenter(i), hist_jet_ET->GetBinCenter(i),
                hist_jet_mass->GetBinCenter(i), hist_jet_DeltaEta->GetBinCenter(i),
                hist_jet_Deltaphi->GetBinCenter(i), hist_jet_bTag->GetBinContent(i),
                hist_jet_EhadOverEem->GetBinCenter(i), hist_missingET_MET->GetBinCenter(i),
                hist_missingEt_Phi->GetBinCenter(i), hist_scalarHT_HT->GetBinCenter(i)
            };

            

            Double_t weight = hist_b1_pt->GetBinContent(i); // Use bin content as event weight

            if (fileName.find("signal") != std::string::npos) {
                loader.AddSignalTrainingEvent(values, weight);
                loader.AddSignalTestEvent(values, weight);
            } else {
                loader.AddBackgroundTrainingEvent(values, weight);
                loader.AddBackgroundTestEvent(values, weight);
            }
        }
    }

    // Prepare training and test trees
    std::cout << "Preparing training and test trees..." << std::endl;
    loader.PrepareTrainingAndTestTree("", "", "nTrain_Signal=1000:nTrain_Background=1000:SplitMode=Random:NormMode=NumEvents:!V");

    // Configure and book BDT
    std::cout << "Booking BDT method..." << std::endl;
    factory.BookMethod(&loader, TMVA::Types::kBDT, "BDT",
                       "!H:!V:NTrees=850:MinNodeSize=2.5%:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:SeparationType=GiniIndex:nCuts=20");

    // Train, test, and evaluate
    std::cout << "Training BDT..." << std::endl;
    factory.TrainAllMethods();

    std::cout << "Testing BDT..." << std::endl;
    factory.TestAllMethods();

    std::cout << "Evaluating BDT..." << std::endl;
    factory.EvaluateAllMethods();

    std::cout << "Closing output file..." << std::endl;
    outputFile->Close();
     
     
     
     
}