#include <stdio.h>
#include <iostream>
#include <map>
#include <string>

#include "TChain.h"
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TObjString.h"
#include "TSystem.h"
#include "TROOT.h"

#include "TMVA/Factory.h"
#include "TMVA/DataLoader.h"
#include "TMVA/Tools.h"
#include "TMVA/TMVAGui.h"


int TMVAClassification(TString myMethodList="")
{
    TMVA::Tools::Instance();

    TFile *signalfile = new TFile("SingleElectron.root","OPEN");
    TFile *datafile = new TFile("Background.root","OPEN");

    TTree *signalTree = (TTree*)signalfile->Get("Signal_Tree");
    TTree *dataTree = (TTree*)datafile->Get("Bkg_Tree");

    TString outfileName("output.root");
    TFile* outputFile = TFile::Open(outfileName,"RECREATE");

    //TMVA::Factory *factory = new TMVA::Factory("TMVAClassification_svm",outputFile,"!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification");
    TMVA::Factory *factory = new TMVA::Factory("TMVAClassification",outputFile,"!V:!Silent:Color:DrawProgressBar:Transformations=G:AnalysisType=Classification");

    TMVA::DataLoader *dataloader=new TMVA::DataLoader("dataset");

    dataloader->AddVariable("ntEgEt","ntEgEt", "units", 'F');
    dataloader->AddVariable("ntEgEta", "ntEgEta", "units", 'F');
    dataloader->AddVariable("ntEgPhi", "ntEgPhi", "units", 'F');
    dataloader->AddVariable("NumberOfTracks", "NumberOfTracks", "units", 'i');
    dataloader->AddVariable("Isolation_value", "Isolation_value", "units", 'F');
    dataloader->AddVariable("Track_pT", "Track_pT", "units", 'F');
    dataloader->AddVariable("track_Eta", "track_Eta", "units", 'F');
    dataloader->AddVariable("track_Phi", "track_Phi", "units", 'F');
    dataloader->AddVariable("track_R", "track_R", "units", 'F');

    Double_t signalWeight = 1.0;
    Double_t backgroundWeight = 1.0;

    dataloader->AddSignalTree(signalTree,signalWeight);
    dataloader->AddBackgroundTree(dataTree,backgroundWeight);

    //dataloader->SetBackgroundWeightExpression("weight");

    TCut mycuts = "";
    TCut mycutb = "";


    dataloader->PrepareTrainingAndTestTree(mycuts,mycutb,"nTrain_Signal=1000:nTrain_Background=1000:SplitMode=Random:NormMode=NumEvents:!V");

    //Support Vector Machine
    factory->BookMethod(dataloader,TMVA::Types::kSVM,"SVM","Gamma=0.25:Tol=0.001:VarTransform=Norm");

    //Boosted Decision Tree
    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");

    //Artificial Neural Network - Multi Layer Perceptron
    factory->BookMethod(dataloader,TMVA::Types::kMLP,"MLP","H:!V:NeuronType=tanh:VarTransform=N:NCycles=600:HiddenLayers=N+5:TestRate=5:!UseRegulator");


    factory->TrainAllMethods();
    factory->TestAllMethods();
    factory->EvaluateAllMethods();

    outputFile->Close();

    std::cout <<"==> Wrote root file: " << outputFile->GetName() << std::endl;
    std::cout <<"==> TMVAClassification is done!" << std::endl;

    delete factory;
    delete dataloader;

    if(!gROOT->IsBatch()) TMVA::TMVAGui(outfileName);

    return 0;
}