//
//  PCA_class_root6.cpp
//  root
//
//  Created by Anja Novosel on 21/05/25.
//
//  Amended based on tau to electron Belle and current MC (removed cont_nbout) - Andre
//

#include <TFile.h>
#include <TTree.h>
#include <TCanvas.h>
#include <TH1F.h>
#include <TLegend.h>
#include <TStyle.h>
#include <TSystem.h>
#include <TString.h>
#include "TMVA/Factory.h"
#include "TMVA/DataLoader.h"
#include "TMVA/Tools.h"
#include "TMVA/TMVAGui.h"
#include <vector>
#include <string>
#include <iostream>

void Bzptn_e_belle_mva(const char* filename, const char* BDT_name1, const char* BDT_name2) {
    
    // Activate TMVA
    TMVA::Tools::Instance();
    
    // to get access to the GUI and all tmva macros
    //TString tmva_dir(TString(gRootDir) + "/tmva");
    //if(gSystem->Getenv("TMVASYS"))
    //    tmva_dir = TString(gSystem->Getenv("TMVASYS"));
    //gROOT->SetMacroPath(tmva_dir + "/test/:" + gROOT->GetMacroPath() );
    //gROOT->ProcessLine(".L TMVAGui.C");


    // Open input file
    TFile* inputFile_train = TFile::Open("/import/wu1/ahua8340/phd/pi_tau_nu_MC/SLptn_e_mix_sig_train_belle.root");
    TFile* inputFile_test = TFile::Open("/import/wu1/ahua8340/phd/pi_tau_nu_MC/SLptn_e_mix_sig_test_belle.root");
    
    // Output file for TMVA
    TString outfileName(filename);
    TFile* outputFile = TFile::Open( outfileName, "RECREATE" );
    

    
        
    // Create TMVA factory and dataloader
    TMVA::Factory *factory = new TMVA::Factory("TMVAClassification", outputFile,
                                               "!V:!Silent:Color:DrawProgressBar:Transformations=I;D:AnalysisType=Classification");
    
    auto dataloader = new TMVA::DataLoader("Belle_e_tmva");   
    // Add variables

    dataloader->AddVariable("weMissM2__boY4ScleanMask__cm__sp0__bc", 'F');		// Missing Mass Squared
    dataloader->AddVariable("lep_CMS_p", 'F');						// Lepton momentum
    dataloader->AddVariable("pi_CMS_p", 'F');						// Pion momentum
    //dataloader->AddVariable("m5", 'F');						// m5
    dataloader->AddVariable("m11", 'F');						// m11
    dataloader->AddVariable("m15", 'F');						// m15
    dataloader->AddVariable("m59",'F');							// m59
    dataloader->AddVariable("m510",'F');						// m510
    //dataloader->AddVariable("Bsig_recQ2Bh", "q^2", "", 'F');				// q^2
    
    // Spectator variables
    dataloader->AddSpectator("roeEextra__boY4ScleanMask__bc",'F');			// ROE E_ECL 


        
    //TTree *inputTrain = (TTree*)inputFile_train->Get("variables");
    //TTree *inputTest = (TTree*)inputFile_test->Get("variables");    
   
    TCut signalcut = "Signal==1";
    TCut backgroundcut = "Signal==0";
    
    Double_t sigWeight = 1.0;
    Double_t bkgWeight = 1.0;
    
    dataloader->AddTree(inputFile_train->Get<TTree>("variables"), "Signal", sigWeight, signalcut, TMVA::Types::kTraining);
    dataloader->AddTree(inputFile_train->Get<TTree>("variables"), "Background", bkgWeight, backgroundcut, TMVA::Types::kTraining);    
    dataloader->AddTree(inputFile_test->Get<TTree>("variables"), "Signal", sigWeight, signalcut, TMVA::Types::kTesting);
    dataloader->AddTree(inputFile_test->Get<TTree>("variables"), "Background", bkgWeight, backgroundcut, TMVA::Types::kTesting);
        
    //TTree *signal     = (TTree*)inputFile_sig->Get("variables");
    //TTree *background = (TTree*)inputFile_bkg->Get("variables");
    
    //factory->AddSignalTree    ( signal);
    //factory->AddBackgroundTree( background);
    
    // Apply cuts on the signal and BKG events. Filter electron events.
    //TCut mycuts = "weMissM2__boY4ScleanMask__cm__sp0__bc>-0.5 && roeEextra__boY4ScleanMask__bc<1.0 && tau_dmID==1";
    //TCut mycutb =  "weMissM2__boY4ScleanMask__cm__sp0__bc>-0.5 && roeEextra__boY4ScleanMask__bc<1.0 && tau_dmID==1";

    // Apply cuts on the signal and BKG events: Take all available events and use 1/2 for train and 1/2 for test.
    //factory->PrepareTrainingAndTestTree( mycuts, mycutb, "nTrain_Signal=0:nTrain_Background=0:SplitMode=Alternate:NormMode=NumEvents:!V");
    
    // Book BDT method
    // No Variable Transform
    factory->BookMethod(dataloader, TMVA::Types::kBDT, BDT_name1, "!H:!V:CreateMVAPdfs:NTrees=400:BoostType=Grad:UseBaggedBoost:BaggedSampleFraction=0.5:MaxDepth=5");
    // With Decorrelation
    factory->BookMethod(dataloader, TMVA::Types::kBDT, BDT_name2, "!H:!V:CreateMVAPdfs:NTrees=400:BoostType=Grad:UseBaggedBoost:BaggedSampleFraction=0.5:MaxDepth=5:VarTransform=D");
    
    // Train, test, evaluate
    factory->TrainAllMethods();
    factory->TestAllMethods();
    factory->EvaluateAllMethods();
    
    // Close output
    outputFile->Close();
    
    delete factory;
    

    std::cout << "TMVA BDT training complete.\n" << std::endl;
    
   if (!gROOT->IsBatch()) TMVA::TMVAGui( outfileName );
 }

int Bzptn_e_belle_tmva() {

   Bzptn_e_belle_mva("/suphys/ahua8340/phd_analysis/belle_mva/tmva/Bzptn_e_belle_mva.root", "Belle_e_id", "Belle_e_de");

   return 0;
}