#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"



void TMVAnalysis(int nvar, float signal, float bk, char output[]) {

    TFile* outputFile = TFile::Open( output, "RECREATE" );
    TFile* inputBKFile = TFile::Open( "/home/ubuntu/samples/original/ttbar-Madgraph-MLM.root", "" );
    TFile* inputSIGNALFile = TFile::Open( "/home/ubuntu/samples/original/ttdm-scalar_100_1.root", "" );

    TTree* BKtree = (TTree *)inputBKFile->Get("t");
    TTree* SIGNALtree = (TTree *)inputSIGNALFile->Get("t");
    BKtree->SetEntries(10000);
    SIGNALtree->SetEntries(10000);

    TMVA::Factory *factory = new TMVA::Factory( "MVAnalysis", outputFile,"!V");
    TMVA::DataLoader *dataloader=new TMVA::DataLoader("datasetBkg0");
   
    dataloader->AddSignalTree ( SIGNALtree, signal );
    dataloader->AddBackgroundTree ( BKtree, bk );

	//variables por orden de "relevancia", se eligen las nvar primeras
    char vars[16][8]={"mMET", "ptMET", "etaMET", "ptb1", "ptb2", "ptlep1", "ptlep2", "etab1", "etalep1", "etab2", "etalep2", "philep2", "phib1", "phib2", "phiMET", "philep1"};
    for (int i=0; i<nvar; i++) {
		dataloader->AddVariable(vars[i],'F');
	}

    dataloader->PrepareTrainingAndTestTree("", "SplitMode=Random:!V" );

    //Neural network (Multilayer Perceptron)
    factory->BookMethod(dataloader, TMVA::Types::kMLP, "MLP8-4", "!V:NCycles=400:HiddenLayers=8,4:TestRate=5:VarTransform=N");
    factory->BookMethod(dataloader, TMVA::Types::kMLP, "MLP8", "!V:NCycles=400:HiddenLayers=8:TestRate=5:VarTransform=N");
    factory->BookMethod(dataloader, TMVA::Types::kMLP, "MLP12-4", "!V:NCycles=200:HiddenLayers=12,4:TestRate=5:VarTransform=N");
    //BDT
    factory->BookMethod(dataloader, TMVA::Types::kBDT, "BDT", "!H:!V:NTrees=800:MinNodeSize=2.5%:MaxDepth=3:BoostType=AdaBoost:AdaBoostBeta=0.2:UseBaggedBoost:BaggedSampleFraction=0.5:SeparationType=GiniIndex:nCuts=20:VarTransform=N" );
    factory->BookMethod(dataloader, TMVA::Types::kCuts, "Cuts", "VarTransform=N");
    factory->BookMethod(dataloader, TMVA::Types::kLD,"LD", "VarTransform=N");
    factory->BookMethod(dataloader, TMVA::Types::kSVM,"SVM", "VarTransform=N");
    factory->BookMethod(dataloader,TMVA::Types::kFisher,"Fisher", "VarTransform=N");

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

    delete factory;
}