#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
#include <TFile.h>
#include <TTree.h>
#include <TH1D.h>
#include "Event.h"
#include <TMVA/Reader.h>

using namespace std;

int main() {

// Set up an output file and book some histograms

  TFile* histFile = new TFile("analysis.root", "RECREATE");
  TH1D* hFishSig = new TH1D("hFishSig", "Fisher, signal", 100, -10.0, 10.0);
  TH1D* hFishBkg = new TH1D("hFishBkg", "Fisher, background", 100, -10.0, 10.0);

// Set up the TMVA Reader object.
// The names in AddVariable must be same as in the input (weight) file.

  TMVA::Reader* reader = new TMVA::Reader();
  float x, y, z;                       // TMVA needs float, not double
  reader->AddVariable("x", &x);  
  reader->AddVariable("y", &y);
  reader->AddVariable("z", &z);
  string dir    = "../train/weights/";
  string prefix = "tmvaTest";
  reader->BookMVA("Fisher", dir + prefix + "_Fisher.weights.xml");

// Open input file, get the TTrees, put into a vector

  TFile* inputFile = new TFile("../generate/testData.root");
  // inputFile->ls();

  TTree* sig = dynamic_cast<TTree*>(inputFile->Get("sig"));
  TTree* bkg = dynamic_cast<TTree*>(inputFile->Get("bkg"));
  vector<TTree*> treeVec;
  treeVec.push_back(sig);
  treeVec.push_back(bkg);

  // loop over cut values

  double tCutMin = -4.;
  double tCutMax = 4.;
  int numCutVal = 101;

  for (int k=0; k<numCutVal; k++){

    double tCut = (tCutMax - tCutMin)*static_cast<double>(k) / 
             static_cast<double>(numCutVal - 1)  + tCutMin;  

// Loop over TTrees (sig and bkg)

    int nSigAccFish = 0;
    int nBkgAccFish = 0;
    int nSig, nBkg;
    for (int i=0; i<treeVec.size(); i++){

      // treeVec[i]->Print();
      Event evt;
      treeVec[i]->SetBranchAddress("evt", &evt);
      int numEntries = treeVec[i]->GetEntries();
      if ( i == 0 ) { nSig = numEntries; }
      if ( i == 1 ) { nBkg = numEntries; }
      // cout << "number of entries = " << numEntries << endl;

// Loop over events.  The test statistic is identified by the first 
// argument used above in BookMVA (below, e.g., "Fisher")

      for (int j=0; j<numEntries; j++){
        treeVec[i]->GetEntry(j);                // sets evt
        x = evt.x;                              // set variables of reader
        y = evt.y;
        z = evt.z;
        double tFisher = reader->EvaluateMVA("Fisher");
        if ( i == 0 ){                       // signal
          hFishSig->Fill(tFisher);
          if ( tFisher >= tCut ) { nSigAccFish++; }
        }
        else if ( i == 1 ){                  // background
          hFishBkg->Fill(tFisher);
          if ( tFisher >= tCut ) { nBkgAccFish++; }
        }

      }

    }

    double epsSigFish = static_cast<double>(nSigAccFish)/
                        static_cast<double>(nSig);
    double epsBkgFish = static_cast<double>(nBkgAccFish)/
                        static_cast<double>(nBkg);

       // cout << "Fisher signal efficiency     = " << epsSigFish << endl;
       // cout << "Fisher background efficiency = " << epsBkgFish << endl;

  // Discovery significance

    double lumi = 20.;         // fb-1
    double sigma_s = 0.2;       // fb
    double sigma_b = 10.;      // fb

    double s = epsSigFish * sigma_s * lumi;
    double b = epsBkgFish * sigma_b * lumi;
    double Z = 0.;
    if ( b > 0. ) { Z = s / sqrt(b); }
    double Z_A = 0.;
    if ( b > 0. ) { Z_A = sqrt( 2. * ( (s + b) * log(1. + s/b) - s) ); }

       //   cout << "s = " << s << endl;
       // cout << "b = " << b << endl;
       // cout << "Z = " << Z << endl;

    cout << tCut << "  " << s << "  " << b << "  " << Z << "  " << Z_A << endl;

  }
  
  histFile->Write();
  histFile->Close();
  return 0;

}