#include <iostream>
#include <TFile.h>
#include <TTree.h>
#include <TMVA/Factory.h>
#include <TMVA/DataLoader.h>
#include <TMVA/Tools.h>

void SetupTMVA() {
    // Initialization
    TFile* input = nullptr;
    TString fname[1] = {
        //"/home/ahmedsayed/TMVA/h90p1_tree.root",
       // "/home/ahmedsayed/TMVA/h90p2_tree.root",
        //"/home/ahmedsayed/TMVA/h95p1_tree.root",
        //"/home/ahmedsayed/TMVA/h95p31_tree.root",
        //"/home/ahmedsayed/TMVA/h96p1_tree.root",
        //"/home/ahmedsayed/TMVA/h96p2_tree.root",
        //"/home/ahmedsayed/TMVA/h96p3_tree.root",
       // "/home/ahmedsayed/TMVA/h96p6_tree.root",
        //"/home/ahmedsayed/TMVA/h_tree.root",
        //"/home/ahmedsayed/TMVA/z_tree.root",
        "/home/ahmedsayed/TMVA/tree.root"
    };

    TFile::SetCacheFileDir(".");
    input = TFile::Open(fname[0], "CACHEREAD");
    
    if (!input) {
        std::cerr << "ERROR: Could not open data file" << std::endl;
        exit(1);
    }
    
    std::cout << "--- TMVAClassification: Using input file: " << input->GetName() << std::endl;
    
    // Register the training and test trees
  /*  TTree* sig90p1Tree = (TTree*)input->Get("Delphes");
    TTree* sig90p2Tree = (TTree*)input->Get("Delphes");
    TTree* sig95p1Tree = (TTree*)input->Get("Delphes");
    TTree* sig95p3Tree = (TTree*)input->Get("Delphes");
    TTree* sig96p1Tree = (TTree*)input->Get("Delphes");
    TTree* sig96p2Tree = (TTree*)input->Get("Delphes");
    TTree* sig96p3Tree = (TTree*)input->Get("Delphes");
    TTree* sig96p6Tree = (TTree*)input->Get("Delphes");
    
    TTree* bkghTree = (TTree*)input->Get("Delphes");
    TTree* bkgzTree = (TTree*)input->Get("Delphes");
    */
    TTree* bkgBBTree = (TTree*)input->Get("Delphes");
    
    // Create a ROOT output file where TMVA will store ntuples, histograms, etc.
    TString outfileName("TMVA.root");
    TFile* outputFile = TFile::Open(outfileName, "RECREATE");
    
    // Create the factory object
    TMVA::Factory* factory = new TMVA::Factory("TMVAClassification", outputFile,
                                               "!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification");
    
    TMVA::DataLoader* dataloader = new TMVA::DataLoader("dataset");
    
    // Define the input variables
    dataloader->AddVariable("b1_pt", "B1quark_Pt", "GEV/c", 'F');
    dataloader->AddVariable("b2_pt", "B2quark_Pt", "GEV/c", 'F');
    dataloader->AddVariable("b1_eta", "B1quark_Eta", "units", 'F');
    dataloader->AddVariable("b2_eta", "B2quark_Eta", "units", 'F');
    dataloader->AddVariable("h_mass", "Higgs_Mass", "GEV", 'F');
    dataloader->AddVariable("h_pt", "Higgs_Pt", "GEV/c", 'F');
    dataloader->AddVariable("h_eta", "Higgs_Eta", "units", 'F');
    dataloader->AddVariable("scalarHt_Ht", "ScalarHt_Ht", "GEV/c", 'F');
    
    
    
    // Define weights
  //  Double_t signalWeights[8] = {0.86430, 0.81540, 1.18770, 0.762, 1.23660, 1.25430, 0.7920, 1.2093};
    //Double_t backgroundWeights[3] = {0.0766, 2189.1, 1811700};
    Double_t backgroundWeights[1] ={1811700};
    // Add signal and background trees
  /*  dataloader->AddSignalTree(sig90p1Tree, signalWeights[0]);
    dataloader->AddSignalTree(sig90p2Tree, signalWeights[1]);
    dataloader->AddSignalTree(sig95p1Tree, signalWeights[2]);
    dataloader->AddSignalTree(sig95p3Tree, signalWeights[3]);
    dataloader->AddSignalTree(sig96p1Tree, signalWeights[4]);
    dataloader->AddSignalTree(sig96p2Tree, signalWeights[5]);
    dataloader->AddSignalTree(sig96p3Tree, signalWeights[6]);
    dataloader->AddSignalTree(sig96p6Tree, signalWeights[7]);
    
    dataloader->AddBackgroundTree(bkghTree, backgroundWeights[0]);
    dataloader->AddBackgroundTree(bkgzTree, backgroundWeights[1]); 
    */
    dataloader->AddBackgroundTree(bkgBBTree, backgroundWeights[0]);
    
    // Set event weights
   // dataloader->SetBackgroundWeightExpression("weight");
    
    // Apply additional cuts
    TCut mycut = "abs(b1_pt > 40 && b2_pt > 30 && h_mass > 40 && h_pt > 40 && h_eta < 10 && scalarHt_Ht > 60)";
    
    // Prepare training and test trees
    dataloader->PrepareTrainingAndTestTree(mycut,
                                           "nTrain_Signal=300000:nTrain_Background=300000:SplitMode=Random:NormMode=NumEvents:!V");
    
    // Book MVA methods
    // Example for booking methods, replace 'Use' conditions with actual flags or logic
    bool Use[] = {true,true,true,true }; // Modify according to actual usage flags
    
    if (Use[0]) {
        factory->BookMethod(dataloader, TMVA::Types::kCuts, "Cuts",
                            "!H:!V:FitMethod=MC:EffSel:SampleSize=200000:VarProp=FSmart");
    }
    
    if (Use[1]) {
        factory->BookMethod(dataloader, TMVA::Types::kLikelihood, "Likelihood",
                            "H:!V:TransformOutput:PDFInterpol=Spline2:NSmoothSig[0]=20:NSmoothBkg[0]=20:NSmoothBkg[1]=10:NSmooth=1:NAvEvtPerBin=50");
    }
    
    if (Use[2]) {
        factory->BookMethod(dataloader, TMVA::Types::kMLP, "MLP",
                            "H:!V:NeuronType=tanh:VarTransform=N:NCycles=600:HiddenLayers=N+5:TestRate=5:!UseRegulator");
    }
    
    if(Use[3]){  
		
		factory->BookMethod( dataloader, TMVA::Types::kBDT, "BDT",
                            "!H:!V:NTrees=850:MinNodeSize=2.5%:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.5:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20" );
 		
		}
    

	
		
    
    
    // Train and evaluate methods
    factory->TrainAllMethods();
    factory->TestAllMethods();
    factory->EvaluateAllMethods();
    
    // Cleanup
    delete factory;
    delete dataloader;
    outputFile->Close();
    delete outputFile;
    input->Close();
    delete input;


}