#include <cstdlib>
#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/Tools.h"
#include "TMVA/TMVAGui.h"
#include <TMVA/Config.h>

using namespace std;

int TMVAClassification( TString myMethodList = "" )
{

   // loads the library
   TMVA::Tools::Instance();

   // default MVA methods to be trained + tested
   std::map<std::string,int> Use;
   // --- Boosted Decision Trees
   Use["BDT"]             = 1; // uses Adaptive Boost
   Use["BDTG"]            = 0; // uses Gradient Boost
   Use["BDTB"]            = 0; // uses Bagging
   Use["BDTD"]            = 0; // decorrelation + Adaptive Boost
   Use["BDTF"]            = 0; // allow usage of fisher discriminant for node splitting 

   std::cout << "\n==> Start TMVAClassification" << std::endl;

   // select methods 
   if (myMethodList != "") {
      for (std::map<std::string,int>::iterator it = Use.begin(); it != Use.end(); it++) it->second = 0;
      std::vector<TString> mlist = TMVA::gTools().SplitString( myMethodList, ',' );
      for (UInt_t i=0; i<mlist.size(); i++) {
         std::string regMethod(mlist[i]);
         if (Use.find(regMethod) == Use.end()) {
            std::cout << "Method \"" << regMethod << "\" not known in TMVA under this name. Choose among the following:" << std::endl;
            for (std::map<std::string,int>::iterator it = Use.begin(); it != Use.end(); it++) std::cout << it->first << " ";
            std::cout << std::endl;
            return 1;
         }
         Use[regMethod] = 1;
      }
   }


   // root file
   string path{"/.../ntuple/"};

   vector<string> vinput {
       "ntuple.root",
       "..."
   };  

   // input file tree directories
   vector<string>vdir{"mt"};

   for(auto &dir : vdir ){

     // output file 
     TString outfileName{dir+"/"+"TMVA.root"};
     TFile* outputFile = TFile::Open( outfileName, "RECREATE" );

     // create the factory object.
     // TMVAClassification: base of the name of the weightfiles 
     // all output can be suppressed by removing the "!" 
     TMVA::Factory *factory = new TMVA::Factory( "TMVAClassification", outputFile,
	 "!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification" );

     // weight file directory
     (TMVA::gConfig().GetIONames()).fWeightFileDir = dir+"/weights";

     factory->AddVariable( "pt_1",   "pt_1",  "GeV", 'F' ); 
     ...

     // input file
     for (auto &inp : vinput){ 	

       string input{path+inp};

       TFile* ifile = new TFile(input.c_str(),"READ");
       if (!ifile->IsOpen()) { cout<<"TMVAClassification(): can't open "<<ifile->GetName()<<endl; return 1; }

       string treeName{"h2taus/h2taus/"+dir};
       TTree *tree = (TTree*)ifile->Get(treeName.c_str());
       if(!tree) { cout<<"tree "<<tree->GetName()<<" does not exist!!"<<endl; return 1;}
 
       factory->AddSignalTree(tree, weight);
       factory->AddBackgroundTree(tree, weight);
     } 

     // for training, and the other half for testing:
     //factory->PrepareTrainingAndTestTree( mycut, "SplitMode=random:!V" );
     TCut mycut;
     mycut="";
     factory->PrepareTrainingAndTestTree( mycut, "SplitMode=Random:NormMode=NumEvents:!V" ); 

     if (Use["BDTD"]) // Decorrelation + Adaptive Boost
       factory->BookMethod( TMVA::Types::kBDT, "BDTD",
     	   "!H:!V:NTrees=400:MinNodeSize=5%:MaxDepth=3:BoostType=AdaBoost:SeparationType=GiniIndex:nCuts=20:VarTransform=Decorrelate" );

     // train MVAs using the set of training events
     factory->TrainAllMethods();
     // evaluate all MVAs using the set of test events
     factory->TestAllMethods();
     // evaluate and compare performance of all configured MVAs
     factory->EvaluateAllMethods();

     // Save the output
     outputFile->Close();
     std::cout << "TMVAClassification(): wrote output: " << outputFile->GetName() << std::endl;

     delete factory;
   } 
   return 0;
} 


int main( int argc, char** argv )
{
  // Select methods 
  TString methodList; 
  for (int i=1; i<argc; i++) {
    TString regMethod(argv[i]);
    if(regMethod=="-b" || regMethod=="--batch") continue;
    if (!methodList.IsNull()) methodList += TString(","); 
    methodList += regMethod;
  }
  return TMVAClassification(methodList); 
}