#include <cstdlib>
#include <iostream>
#include <map>
#include <string>
#include <fstream>
#include <glob.h>

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

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

int TMVACrossValidation_Exercise_TTFCNC(bool useRandomSplitting = false)
{
   (TMVA::gConfig().GetVariablePlotting()).fMaxNumOfAllowedVariablesForScatterPlots = 41;

   TMVA::Tools::Instance();

   std::map<std::string,int> Use;

#ifdef R__HAS_TMVAGPU
   Use["DNN_GPU"]         = 1;
#else
   Use["DNN_GPU"]         = 0;
#endif

#ifdef R__HAS_TMVACPU
   Use["DNN_CPU"]         = 1;
#else
   Use["DNN_CPU"]         = 0;
#endif

   TChain *inputS = new TChain("Events");
   TChain *inputB = new TChain("Events");

   std::string sig_in_line;
   ifstream sig_in("Automation_Cla/filelist_sig_ttzct.txt");
   while (std::getline(sig_in, sig_in_line)){
      TString sig_in_line_R = sig_in_line; 
      inputS->Add(sig_in_line_R);
   }

   std::string bkg_in_line;
   ifstream bkg_in("Automation_Cla/filelist_bkg.txt");
   while (std::getline(bkg_in, bkg_in_line)){
      TString bkg_in_line_R = bkg_in_line;
      inputB->Add(bkg_in_line_R);
   }

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

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

   dataloader->AddVariable( "KinTopWb_eta", "SMTop eta", 'F' );
   dataloader->AddVariable( "KinTopWb_mass", "SMTop mass", 'F' );
   dataloader->AddVariable( "KinTopWb_phi", "SMTop phi", 'F' );
   dataloader->AddVariable( "KinTopZq_eta", "FCNCTop eta", 'F' );
   dataloader->AddVariable( "KinTopZq_mass", "FCNCTop mass", 'F' );
   dataloader->AddVariable( "KinTopZq_phi", "FCNCTop phi", 'F' );
   dataloader->AddVariable( "KinTopZq_pt", "FCNCTop transverse momentum", 'F' );
   dataloader->AddVariable( "MET_phi", "MET_Phi", 'F' );
   dataloader->AddVariable( "MET_pt", "MET", 'F' );
   dataloader->AddVariable( "MVAinput_bJ_DeepJetB", "DeepJetBtagger of SM jet", 'F' );
   dataloader->AddVariable( "MVAinput_bJ_pt", "SM jet transverse momentum", 'F' );
   dataloader->AddVariable( "MVAinput_bJqJ_dR", "dR of SM/FCNC jet", 'F' );
   dataloader->AddVariable( "MVAinput_qJ_DeepJetB", "DeepJetBtagger of FCNC jet", 'F' );
   dataloader->AddVariable( "MVAinput_qJ_pt", "FCNC jet transverse momentum", 'F' );
   dataloader->AddVariable( "MVAinput_WLbJ_dR", "dR of WL/SM jet", 'F' );
   dataloader->AddVariable( "MVAinput_Wlepchaeta", "lepton charge from W * eta", 'F' );
   dataloader->AddVariable( "MVAinput_WLqJ_dR", "dR of WL/FCNC jet", 'F' );
   dataloader->AddVariable( "MVAinput_WLZL2_dR", "dR of WL/ZL2", 'F' );
   dataloader->AddVariable( "MVAinput_ZL1bJ_dR", "dR of ZL1/SM jet", 'F' );
   dataloader->AddVariable( "MVAinput_ZL1qJ_dR", "dR of ZL1/FCNC jet", 'F' );
   dataloader->AddVariable( "MVAinput_ZL1ZL2_dR", "dR of ZL1/ZL2", 'F' );
   dataloader->AddVariable( "MVAinput_ZL2bJ_dR", "dR of ZL2/SM jet", 'F' );
   dataloader->AddVariable( "MVAinput_ZL2qJ_dR", "dR of ZL2/FCNC jet", 'F' );
   dataloader->AddVariable( "TriLepton_mass", "Trilepton mass", 'F' );
   dataloader->AddVariable( "TriLepton_WleptonZdPhi", "dPhi of Z/WL", 'F' );
   dataloader->AddVariable( "TriLepton_WleptonZdR", "dR of Z/WL", 'F' );
   dataloader->AddVariable( "W_MT", "W mass", 'F' );
   dataloader->AddVariable( "Z_mass", "Z mass", 'F' );

   dataloader->AddSpectator("eventID", 'I');

   Double_t signalWeight     = 1.0;
   Double_t backgroundWeight = 1.0;

   dataloader->AddSignalTree    ( inputS ,     signalWeight );
   dataloader->AddBackgroundTree( inputB , backgroundWeight );

   dataloader->SetSignalWeightExpression( "(abs(genWeight)/genWeight)*puWeight*JetPUID_SF*Electron_SF*Muon_SF*BtagWeight" );
   dataloader->SetBackgroundWeightExpression( "(abs(genWeight)/genWeight)*puWeight*JetPUID_SF*Electron_SF*Muon_SF*BtagWeight" );

   TCut mycuts = "(CutStep>=1)&&(nGoodJet>=2)&&(nGoodJet<=3)&&(nBjet>=1)&&(JetVetoSum<9)&&((TMath::Power(((KinTopWb_mass-172.5)/59.7),2)+TMath::Power(((KinTopZq_mass-172.5)/41.7),2))<1.0)";
   TCut mycutb = "(CutStep>=1)&&(nGoodJet>=2)&&(nGoodJet<=3)&&(nBjet>=1)&&(JetVetoSum<9)&&((TMath::Power(((KinTopWb_mass-172.5)/59.7),2)+TMath::Power(((KinTopZq_mass-172.5)/41.7),2))<1.0)";

   // Replaces PrepareTrainingAndTestTree
   UInt_t  numFolds = 4;
   UInt_t  nSignalEvents     = 432851;
   UInt_t  nBackgroundEvents = 269441;
   UInt_t  partition_nSignalEvents        = UInt_t(Double_t(nSignalEvents) / Double_t(numFolds+1));
   UInt_t  partition_nBackgroundEvents    = UInt_t(Double_t(nBackgroundEvents) / Double_t(numFolds+1));
   TString TS_partition_nSignalEvents     = TString::Format("%u", partition_nSignalEvents);
   TString TS_partition_nBackgroundEvents = TString::Format("%u", partition_nBackgroundEvents);

   TString nTest_Signal = "nTest_Signal=" + TS_partition_nSignalEvents;
   TString nTest_Background = "nTest_Background=" + TS_partition_nBackgroundEvents;
   TString PrepareTrainingAndTestTree_Option = nTest_Signal + ":" + nTest_Background + ":SplitMode=Random" + ":NormMode=NumEvents" + ":!V";
   dataloader->PrepareTrainingAndTestTree(mycuts, mycutb, PrepareTrainingAndTestTree_Option);

   TString analysisType = "Classification";

   TString splitType = (useRandomSplitting) ? "Random" : "Deterministic";
   TString splitExpr = (!useRandomSplitting) ? "int(fabs([eventID]))%int([NumFolds])" : "";
 
   TString cvOptions = Form("!V"
                            ":!Silent"
                            ":ModelPersistence"
                            ":AnalysisType=%s"
                            ":SplitType=%s"
                            ":NumFolds=%i"
                            ":SplitExpr=%s",
                            analysisType.Data(), splitType.Data(), numFolds,
                            splitExpr.Data());
 
   TMVA::CrossValidation cv{"TMVACrossValidation", dataloader, outputFile, cvOptions};

   if (Use["DNN_CPU"] or Use["DNN_GPU"]) {
      TString layoutString ("Layout=RELU|256,RELU|256,TANH|256,RELU|256,RELU|256,LINEAR");

      TString trainingStrategyString = ("TrainingStrategy=LearningRate=1e-3,Momentum=0.9,"
                                        "ConvergenceSteps=20,BatchSize=512,TestRepetitions=1,"
                                        "WeightDecay=1e-4,Regularization=L2,"
                                        "DropConfig=0.0+0.25+0.25+0.25+0.25+0.25");

      TString dnnOptions ("!H:V:ErrorStrategy=CROSSENTROPY:VarTransform=Norm:"
                          "WeightInitialization=XAVIERUNIFORM");

      dnnOptions.Append (":"); dnnOptions.Append (layoutString);
      dnnOptions.Append (":"); dnnOptions.Append (trainingStrategyString);

      if (Use["DNN_GPU"]) {
         TString gpuOptions = dnnOptions + ":Architecture=GPU";
         cv.BookMethod(TMVA::Types::kDL, "DNN_GPU", gpuOptions);
      }
      if (Use["DNN_CPU"]) {
         TString cpuOptions = dnnOptions + ":Architecture=CPU";
         cv.BookMethod(TMVA::Types::kDL, "DNN_CPU", cpuOptions);
      }
   }

   cv.Evaluate();

   size_t iMethod = 0;
   for (auto && result : cv.GetResults()) {
      std::cout << "Summary for method " << cv.GetMethods()[iMethod++].GetValue<TString>("MethodName")
                << std::endl;
      for (UInt_t iFold = 0; iFold<cv.GetNumFolds(); ++iFold) {
         std::cout << "\tFold " << iFold << ": "
                   << "ROC int: " << result.GetROCValues()[iFold]
                   << ", "
                   << "BkgEff@SigEff=0.3: " << result.GetEff30Values()[iFold]
                   << std::endl;
      }
   }

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

   if (!gROOT->IsBatch()) {
      cv.GetResults()[0].DrawAvgROCCurve(kTRUE, "Avg ROC for BDTG");
      cv.GetResults()[0].DrawAvgROCCurve(kTRUE, "Avg ROC for Fisher");
 
      TMVA::TMVAGui(outfileName);
   }
 
   return 0;
}

int main(int argc, char **argv)
{
   TMVACrossValidation_Exercise_TTFCNC();
}