/// \file
/// \ingroup tutorial_tmva
/// \notebook
///  TMVA Classification Example Using a Convolutional Neural Network
///
/// This is an example of using a CNN in TMVA. We do classification using a toy image data set
/// that is generated when running the example macro
///
/// \macro_image
/// \macro_output
/// \macro_code
///
/// \author Lorenzo Moneta

/***

    # TMVA Classification Example Using a Convolutional Neural Network


**/

/// Helper function to create input images data
/// we create a signal and background 2D histograms from 2d gaussians
/// with a location (means in X and Y)  different for each event
/// The difference between signal and background is in the gaussian width.
/// The width for the background gaussian is slightly larger than the signal width by few % values
///
///

#include <cstdlib>
#include <vector>
#include <iostream>
#include <map>
#include <string>

#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TSystem.h"
#include "TROOT.h"
#include "TStopwatch.h"

#include "TMVA/Tools.h"
#include "TMVA/Reader.h"
#include "TMVA/MethodCuts.h"

using namespace TMVA;

void MakeImagesTree(int n, int nh, int nw)
{

   // image size (nh x nw)
   const int ntot = nh * nw;
   const TString fileOutName = TString::Format("images_data_test_%dx%d.root", nh, nw);

   const int nRndmEvts = 10000; // number of events we use to fill each image
   double delta_sigma = 0.1;    // 5% difference in the sigma
   double pixelNoise = 5;

   double sX1 = 3;
   double sY1 = 3;
   double sX2 = sX1 + delta_sigma;
   double sY2 = sY1 - delta_sigma;

   auto h = new TH2D("h", "h", nh, 0, 10, nw, 0, 10);

   auto f1 = new TF2("f1", "xygaus");
   auto f2 = new TF2("f2", "xygaus");
   TTree t("test_tree", "test_tree");

   TFile f(fileOutName, "RECREATE");

   std::vector<float> x1(ntot);
   std::vector<float> x2(ntot);
   std::vector<float> x(ntot);

   // create signal and background trees with a single branch
   // an std::vector<float> of size nh x nw containing the image data

   std::vector<float> *px = &x;

   t.Branch("vars", "std::vector<float>", &px);

   // std::cout << "create tree " << std::endl;

   t.SetDirectory(&f);

   f1->SetParameters(1, 5, sX1, 5, sY1);
   f2->SetParameters(1, 5, sX2, 5, sY2);
   gRandom->SetSeed(0);
   std::cout << "Filling ROOT tree " << std::endl;
   for (int i = 0; i < n; ++i) {
      if (i % 1000 == 0)
         std::cout << "Generating image event ... " << i << std::endl;
      h->Reset();
      // generate random means in range [3,7] to be not too much on the border
      f1->SetParameter(1, gRandom->Uniform(3, 7));
      f1->SetParameter(3, gRandom->Uniform(3, 7));
      f2->SetParameter(1, gRandom->Uniform(3, 7));
      f2->SetParameter(3, gRandom->Uniform(3, 7));

      bool sigEvt = (gRandom->Rndm() < 0.5);

      if (sigEvt)
         h->FillRandom("f1", nRndmEvts);
      else
         h->FillRandom("f2", nRndmEvts);

      for (int k = 0; k < nh; ++k) {
         for (int l = 0; l < nw; ++l) {
            int m = k * nw + l;
            // add some noise in each bin
            x[m] = h->GetBinContent(k + 1, l + 1) + gRandom->Gaus(0, pixelNoise);
         }
      }
      t.Fill();
   }
   t.Write();

   Info("MakeImagesTree", "Tree with images data written to the file %s", f.GetName());
   t.Print();
   f.Close();
}


void TMVA_CNN_ClassificationApplication( TString myMethodList = "" )
{

   //---------------------------------------------------------------
   // This 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["TMVA_DNN_GPU"]             = 1;
   Use["TMVA_CNN_GPU"]             = 1;

   std::cout << std::endl;
   std::cout << "==> Start TMVAClassificationApplication" << std::endl;


   // --------------------------------------------------------------------------------------------------

   // Create the Reader object

   TMVA::Reader *reader = new TMVA::Reader( "!Color:!Silent" );

   // Create a set of variables and declare them to the reader
   // - the variable names MUST corresponds in name and type to those given in the weight file(s) used
   std::vector<float>  vars(16*16);
   for (unsigned int i = 0; i < vars.size(); i++) {
      std::string varName = "vars";//"vars[" + std::to_string(i) + std::string("]");
      reader->AddVariable(varName.c_str() ,&vars[i] );
   }


   // Book the MVA methods

   TString dir    = "dataset/weights/";
   TString prefix = "TMVA_CNN_Classification";

   // Book method(s)
   for (std::map<std::string,int>::iterator it = Use.begin(); it != Use.end(); it++) {
      if (it->second) {
         TString methodName = TString(it->first) + TString(" method");
         TString weightfile = dir + prefix + TString("_") + TString(it->first) + TString(".weights.xml");
         reader->BookMVA( methodName, weightfile );
      }
   }

   // Book output histograms
   UInt_t nbin = 100;
   TH1F *histBdt(0);
   TH1F *histDnnGpu(0);
   TH1F *histCnnGpu(0);

   if (Use["TMVA_DNN_GPU"]) histDnnGpu = new TH1F("MVA_DNN_GPU", "MVA_DNN_GPU", nbin, -0.1, 1.1);
   if (Use["TMVA_CNN_GPU"]) histCnnGpu = new TH1F("MVA_CNN_GPU", "MVA_CNN_GPU", nbin, -0.1, 1.1);
   if (Use["BDT"])           histBdt     = new TH1F( "MVA_BDT",           "MVA_BDT",           nbin, -0.8, 0.8 );


   // Prepare input tree (this must be replaced by your data source)
   // in this example, there is a toy tree with signal and one with background events
   // we'll later on use only the "signal" events for the test in this example.
   //
   TFile *input(0);
   //TString fname = "images_data_test_16x16.root";
   TString fname = "images_data_16x16.root";
   // if (gSystem->AccessPathName( fname )) {
   //    MakeImagesTree(10000, 16, 16);
   // }
   input = TFile::Open( fname ); // check if file in local directory exists
   if (!input) {
      std::cout << "ERROR: could not open data file" << std::endl;
      exit(1);
   }
   std::cout << "--- TMVA_CNN_ClassificationApp    : Using input file: " << input->GetName() << std::endl;

   // Event loop

   // Prepare the event tree
   // - Here the variable names have to corresponds to your tree
   // - You can use the same variables as above which is slightly faster,
   //   but of course you can use different ones and copy the values inside the event loop
   //
   std::cout << "--- Select signal sample" << std::endl;
   //TTree* theTree = (TTree*)input->Get("test_tree");
   TTree* theTree = (TTree*)input->Get("sig_tree");
#if 0
   Float_t userVar1, userVar2;
   theTree->SetBranchAddress( "var1", &userVar1 );
   theTree->SetBranchAddress( "var2", &userVar2 );
   theTree->SetBranchAddress( "var3", &var3 );
   theTree->SetBranchAddress( "var4", &var4 );
#endif
   std::vector<float> * userVars = nullptr;
   theTree->SetBranchAddress( "vars", &userVars );

   // Efficiency calculator for cut method

   std::cout << "--- Processing: " << theTree->GetEntries() << " events" << std::endl;
   TStopwatch sw;
   sw.Start();
   for (Long64_t ievt=0; ievt<theTree->GetEntries();ievt++) {

      if (ievt%1000 == 0) std::cout << "--- ... Processing event: " << ievt << std::endl;

      theTree->GetEntry(ievt);

      // Return the MVA outputs and fill into histograms
      //std::cout << (*userVars).size() << "  " << (*userVars)[0] << std::endl;

      // copy read data in vector associated to the Reader
      std::copy(userVars->begin(), userVars->end(), vars.begin());

      if (Use["TMVA_DNN_GPU"]) histDnnGpu->Fill(reader->EvaluateMVA("TMVA_DNN_GPU method"));
      if (Use["TMVA_CNN_GPU"]) histCnnGpu->Fill(reader->EvaluateMVA("TMVA_CNN_GPU method"));
      if (Use["BDT"          ])   histBdt    ->Fill( reader->EvaluateMVA("BDT method"           ) );

   }

   // Get elapsed time
   sw.Stop();
   std::cout << "--- End of event loop: "; sw.Print();
   // Write histograms

   TFile *target  = new TFile( "TMVACNNApp.root","RECREATE" );
   if (Use["TMVA_DNN_GPU"]) histDnnGpu->Write();
   if (Use["TMVA_CNN_GPU"]) histCnnGpu->Write();
   if (Use["BDT"          ])   histBdt    ->Write();

   target->Close();

   std::cout << "--- Created root file: \"TMVACNNApp.root\" containing the MVA output histograms" << std::endl;

   delete reader;

   std::cout << "==> TMVAClassificationApplication is done!" << std::endl << std::endl;
}