#include <cstdlib>
#include <iostream>
#include <map>
#include <string>
#include <vector>
#include "TFile.h"
#include "TTree.h"
#include "TString.h"
#include "TSystem.h"
#include "TROOT.h"
#include "TStopwatch.h"
#include "TH1F.h"
#include "TMVA/Tools.h"
#include "TMVA/Reader.h"

using namespace TMVA;

void pion()
{

   TMVA::Tools::Instance();
   //---------------------------------------------------------------------------------
   // create the Reader object
   TMVA::Reader *reader = new TMVA::Reader( "!Color:!Silent" );
   Float_t var1, var2, var3, var4, var5, var6, var7, var8, var9, var10, var11, var12, var13, var14, var15, var16, var17, var18;

   reader->AddVariable( "mdcp",           &var1 );
   reader->AddVariable( "mdccost",        &var2 );
   reader->AddVariable( "charge",         &var3 );
   reader->AddVariable( "chidedxpi",      &var4 );
   reader->AddVariable( "chidedxk",       &var5 );
   reader->AddVariable( "chidedxp",       &var6 );
   reader->AddVariable( "t11Mod",         &var7 );
   reader->AddVariable( "t12Mod",         &var8 );
   reader->AddVariable( "t21Mod",         &var9 );
   reader->AddVariable( "t22Mod",         &var10);
   reader->AddVariable( "q0Mod",          &var11);
   reader->AddVariable( "energyMod/mdcp",    &var12);
   reader->AddVariable( "eSeedMod/e3x3Mod",  &var13);
   reader->AddVariable( "emcdthetaMod",      &var14);
   reader->AddVariable( "emcnumHitsMod",     &var15);
   reader->AddVariable( "secondMomentMod",   &var16);
   reader->AddVariable( "latMomentMod",      &var17);
   reader->AddVariable( "depthMod",          &var18);

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


   // book the MVA methods
   std::string DNN_file = "./dataset/weights/TMVAMulticlass_DNN.weights.xml";
   reader->BookMVA( "DNN method", DNN_file );

   //---------------------------------------------------------------------------------
   //input file
   TChain *input = new TChain("particle");
   input->Add("./pionTrain.root");
   //---------------------------------------------------------------------------------
   
   // prepare the tree
   int charge, emcnumHits;
   double mdcp, mdccost;
   double t11, t12, t21, t22, q0, depth; 
   double energy, secondMoment, latMoment, eSeed, e3x3, emcdtheta;
   double chidedxpi, chidedxk, chidedxp;

   input->SetBranchAddress( "mdcp",         &mdcp        );
   input->SetBranchAddress( "mdccost",      &mdccost     );
   input->SetBranchAddress( "charge",       &charge      );
   input->SetBranchAddress( "q0Mod",           &q0          );
   input->SetBranchAddress( "t11Mod",          &t11         );
   input->SetBranchAddress( "t12Mod",          &t12         );
   input->SetBranchAddress( "t21Mod",          &t21         );
   input->SetBranchAddress( "t22Mod",          &t22         );
   input->SetBranchAddress( "energyMod",       &energy      );
   input->SetBranchAddress( "eSeedMod",        &eSeed       );
   input->SetBranchAddress( "e3x3Mod",         &e3x3        );
   input->SetBranchAddress( "secondMomentMod", &secondMoment);
   input->SetBranchAddress( "latMomentMod",    &latMoment   );
   input->SetBranchAddress( "emcnumHitsMod",   &emcnumHits  );
   input->SetBranchAddress( "emcdthetaMod",    &emcdtheta   );
   input->SetBranchAddress( "depthMod",        &depth       );
   input->SetBranchAddress( "chidedxpi",    &chidedxpi   );
   input->SetBranchAddress( "chidedxk",     &chidedxk    );
   input->SetBranchAddress( "chidedxp",     &chidedxp    );

   //output
   int charge_;
   double mdcp_, mdccost_;
   double q0_, depth_, energy_, secondMoment_;
   double chidedxpi_, chidedxk_, chidedxp_;
   double dnn[5];
   double eSeedMod, e3x3Mod;

   TFile *file  = new TFile( "./mypion.root","RECREATE" );
   TTree * T = new TTree("particle", "particle");

   T->Branch("charge",         &charge_      );
   T->Branch("mdcp",           &mdcp_        );
   T->Branch("mdccost",        &mdccost_     );
   T->Branch("q0",             &q0_          );
   T->Branch("energy",         &energy_      );
   T->Branch("secondMoment",   &secondMoment_);
   T->Branch("chidedxpi",      &chidedxpi_   );
   T->Branch("chidedxk",       &chidedxk_    );
   T->Branch("chidedxp",       &chidedxp_    );
   T->Branch("probdnn",        dnn,           "dnn[5]/D"           );

   int nevents = input->GetEntries();
   std::cout << nevents << " events need to be process..."<<std::endl;

   int errevent = 0;

   for (Long64_t ievt=0; ievt<nevents;ievt++) {
      if (ievt%10000 == 0){
         std::cout << "--- ... Processing event: " << ievt << std::endl;
      }
      input->GetEntry(ievt);

      var1  = float(mdcp);
      var2  = float(mdccost);
      var3  = float(charge);
      var4  = float(chidedxpi);
      var5  = float(chidedxk);
      var6  = float(chidedxp);
      var7 = float(t11);
      var8 = float(t12); 
      var9 = float(t21); 
      var10= float(t22); 
      var11 = float(q0); 
      var12 = float(energy/mdcp); 
      if( e3x3==0.0 )            var13 = float(0.0); else{ var13 = float(eSeed/e3x3); }
      var14 = float(emcdtheta); 
      var15 = float(emcnumHits); 
      var16 = float(secondMoment); 
      var17 = float(latMoment); 
      var18 = float(depth); 

      charge_       =  charge      ;
      mdcp_         =  mdcp        ;
      mdccost_      =  mdccost     ;
      q0_           =  q0          ;
      energy_       =  energy      ;
      secondMoment_ =  secondMoment;
      chidedxpi_    =  chidedxpi   ;
      chidedxk_     =  chidedxk    ;
      chidedxp_     =  chidedxp    ;

      vector<float> prob = reader->EvaluateMulticlass( "DNN method" );
//cout<<prob[0]<<" "<<prob[1]<<" "<<prob[2]<<endl;

      dnn[0]=0.0;
      dnn[1]=0.0;
      dnn[2]=double(prob[0]);
      dnn[3]=double(prob[1]);
      dnn[4]=double(prob[2]);
    
      if( dnn[2]<=0 && dnn[3]<=0 && dnn[4]<=0) {
         errevent += 1;
         cout<<var1<<" "<<var2<<" "<<var3<<" "<<var4<<" "<<var5<<" "<<var6<<" "<<var7<<" "<<var8<<" "<<var9<<" "<<var10<<" "<<var11<<" "<<var12<<" "<<var13<<" "<<var14<<" "<<var15<<" "<<var16<<" "<<var17<<" "<<var18<<endl;
      }
      T->Fill();

   }

   cout<<errevent<<endl;
   T->Write();
   file->Close();

   delete reader;
   std::cout << "==> TMVAMulticlassApp is done!" << std::endl << std::endl;
}