#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="")
{
    TFile *input(0);

    //test please
    //TString fname = "./gaussian_data_1.root";
    //TString fname = "./gaussian_data_2.root";
    input = TFile::Open(fname);

    TTree *signalTree = (TTree*)input->Get("Sig_tree");
    TTree *backgroundTree = (TTree*)input->Get("Bg_tree");

    TString outfileName("output_svm.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::DataLoader *dataloader=new TMVA::DataLoader("dataset");

    dataloader->AddVariable("var1", "Variable 1", "units", 'F');
    dataloader->AddVariable("var2", "Variable 2", "units", 'F');
    dataloader->AddVariable("var3", "Variable 3", "units", 'F');
    //dataloader->AddVariable("var4", "Variable 4", "units", 'F');
    //dataloader->AddVariable("var5",'F');
    //dataloader->AddVariable("var6",'F');

    Double_t signalWeight = 1.0;
    Double_t backgroundWeight = 1.0;

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

    dataloader->SetBackgroundWeightExpression("weight");

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


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


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


    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;
}