#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 "TMath.h"
#include "TMVA/Reader.h"
#include "TMVA/Tools.h"
#include "dataset/weights/TMVAClassification_DNN_Barrel.class.C"

class ReadDNN_Barrel;

using namespace TMVA;

void validationClass( TString DataFile = "ntu.root")
{

	gStyle->SetOptTitle(0); gStyle->SetOptStat(0);
	gStyle->SetOptFit(0); gStyle->SetStatBorderSize(0);
	gStyle->SetStatX(.49); gStyle->SetStatY(.89);

	TFile *file = new TFile(DataFile);

	float histMin = 0.0;
	float histMax = 1.0;
	int histBins = 1000;

	TH1F *histSgn = new TH1F( "BDT Sgn", "BDT Sgn", histBins, histMin, histMax );
	TH1F *histBkg = new TH1F( "BDT Bkg", "BDT Bkg", histBins, histMin, histMax );

	//READER

	std::vector<std::string> inputVars;

	inputVars.push_back( "muoPt" );
	inputVars.push_back( "abs(muoEta)" );
	inputVars.push_back( "muoSegmComp" );
	inputVars.push_back( "muoChi2LM" );
	inputVars.push_back( "muoChi2LP" );
	inputVars.push_back( "muoGlbTrackTailProb" );
	inputVars.push_back( "muoIValFrac" );
	inputVars.push_back( "muoLWH" );
	inputVars.push_back( "muoTrkKink" );
	inputVars.push_back( "muoGlbKinkFinderLOG" );
	inputVars.push_back( "muoTimeAtIpInOutErr" );
	inputVars.push_back( "muoOuterChi2" );
	inputVars.push_back( "muoInnerChi2" );  
	inputVars.push_back( "muoTrkRelChi2" );
	inputVars.push_back( "muoVMuonHitComb" );
	inputVars.push_back( "muoGlbDeltaEtaPhi" );
	inputVars.push_back( "muoStaRelChi2" );
	inputVars.push_back( "muoTimeAtIpInOut" );
	inputVars.push_back( "muoValPixHits" );
	inputVars.push_back( "muoNTrkVHits" );
	inputVars.push_back( "muoGNchi2" );
	inputVars.push_back( "muoVMuHits" );
	inputVars.push_back( "muoNumMatches" );
	inputVars.push_back( "muoPFiso" );
	inputVars.push_back( "muoQprod" );

	// define classes
	IClassifierReader* classReader = new ReadDNN_Barrel( inputVars );
	cout << "=== Macro        : Class creation was successful" << endl;

	// prepare the tree
	TTree *tree = (TTree*)file ->Get("PDsecondTree");

	vector <int> *v_muoValPixHits = 0, *v_muoNTrkVHits = 0, *v_muoLWH = 0, *v_muoVMuHits = 0, *v_muoVMuonHitComb = 0, *v_muoQprod = 0;
  vector <int> *v_muoLund = 0, *v_muoAncestor = 0, *v_muoType = 0;
  vector <float> *v_muoPt = 0, *v_muoEta = 0, *v_muoPhi = 0;
  vector <float> *v_muoNumMatches = 0, *v_muoChi2LP = 0, *v_muoTrkKink = 0, *v_muoSegmComp = 0, *v_muoChi2LM = 0, *v_muoTrkRelChi2 = 0;
  vector <float> *v_muoGlbTrackTailProb = 0, *v_muoGlbKinkFinderLOG = 0, *v_muoGlbDeltaEtaPhi = 0, *v_muoStaRelChi2 = 0;
  vector <float> *v_muoTimeAtIpInOut = 0, *v_muoTimeAtIpInOutErr = 0, *v_muoInnerChi2 = 0, *v_muoIValFrac = 0, *v_muoGNchi2 = 0, *v_muoOuterChi2 = 0;
  vector <float> *v_muoPFiso = 0;
  vector <int> *v_muoEvt = 0;

	tree->SetBranchAddress( "muoEvt", &v_muoEvt );
	tree->SetBranchAddress( "muoLund", &v_muoLund );
	tree->SetBranchAddress( "muoAncestor", &v_muoAncestor );
	tree->SetBranchAddress( "muoType", &v_muoType );
	tree->SetBranchAddress( "muoPt", &v_muoPt );
	tree->SetBranchAddress( "muoEta", &v_muoEta );
	tree->SetBranchAddress( "muoSegmComp", &v_muoSegmComp );
	tree->SetBranchAddress( "muoChi2LM", &v_muoChi2LM );
	tree->SetBranchAddress( "muoChi2LP", &v_muoChi2LP );
	tree->SetBranchAddress( "muoGlbTrackTailProb", &v_muoGlbTrackTailProb );
	tree->SetBranchAddress( "muoIValFrac", &v_muoIValFrac );
	tree->SetBranchAddress( "muoLWH", &v_muoLWH );
	tree->SetBranchAddress( "muoTrkKink", &v_muoTrkKink );
	tree->SetBranchAddress( "muoGlbKinkFinderLOG", &v_muoGlbKinkFinderLOG );
	tree->SetBranchAddress( "muoTimeAtIpInOutErr", &v_muoTimeAtIpInOutErr );
	tree->SetBranchAddress( "muoOuterChi2", &v_muoOuterChi2 );
	tree->SetBranchAddress( "muoInnerChi2", &v_muoInnerChi2 );  
	tree->SetBranchAddress( "muoTrkRelChi2", &v_muoTrkRelChi2 );
	tree->SetBranchAddress( "muoVMuonHitComb", &v_muoVMuonHitComb );
	tree->SetBranchAddress( "muoNumMatches", &v_muoNumMatches );
	tree->SetBranchAddress( "muoGlbDeltaEtaPhi", &v_muoGlbDeltaEtaPhi );
	tree->SetBranchAddress( "muoStaRelChi2", &v_muoStaRelChi2 );
	tree->SetBranchAddress( "muoTimeAtIpInOut", &v_muoTimeAtIpInOut );
	tree->SetBranchAddress( "muoValPixHits", &v_muoValPixHits );
	tree->SetBranchAddress( "muoNTrkVHits", &v_muoNTrkVHits );
	tree->SetBranchAddress( "muoGNchi2", &v_muoGNchi2 );
	tree->SetBranchAddress( "muoVMuHits", &v_muoVMuHits );
	tree->SetBranchAddress( "muoPFiso", &v_muoPFiso );
	tree->SetBranchAddress( "muoQprod", &v_muoQprod );

	std::vector<double>* inputVec = new std::vector<double>( inputVars.size() );

	//Long64_t nEntries = tree->GetEntries();
	Long64_t nEntries = 10000;

	for (Long64_t ievt=0; ievt<nEntries; ++ievt){

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

		tree->GetEntry(ievt);


		for(int i=0; i<(*v_muoPt).size(); ++i){

			if(!( TMath::Abs((*v_muoLund)[i])!=0 && (*v_muoPt)[i]>2 && (*v_muoAncestor)[i]>=0) ) continue;
			if( !(TMath::Abs((*v_muoEta)[i]) < 1.2)) continue;

			(*inputVec)[0] = (double)(*v_muoPt)[i];
			(*inputVec)[1] = (double)TMath::Abs((*v_muoEta)[i]);
			(*inputVec)[2] = (double)(*v_muoSegmComp)[i];
			(*inputVec)[3] = (double)(*v_muoChi2LM)[i];
			(*inputVec)[4] = (double)(*v_muoChi2LP)[i];
			(*inputVec)[5] = (double)(*v_muoGlbTrackTailProb)[i];
			(*inputVec)[6] = (double)(*v_muoIValFrac)[i];
			(*inputVec)[7] = (double)(*v_muoLWH)[i];
			(*inputVec)[8] = (double)(*v_muoTrkKink)[i];
			(*inputVec)[9] = (double)(*v_muoGlbKinkFinderLOG)[i];
			(*inputVec)[10] = (double)(*v_muoTimeAtIpInOutErr)[i];
			(*inputVec)[11] = (double)(*v_muoOuterChi2)[i];
			(*inputVec)[12] = (double)(*v_muoInnerChi2)[i];  
			(*inputVec)[13] = (double)(*v_muoTrkRelChi2)[i];
			(*inputVec)[14] = (double)(*v_muoVMuonHitComb)[i];
			(*inputVec)[15] = (double)(*v_muoGlbDeltaEtaPhi)[i];
			(*inputVec)[16] = (double)(*v_muoStaRelChi2)[i];
			(*inputVec)[17] = (double)(*v_muoTimeAtIpInOut)[i];
			(*inputVec)[18] = (double)(*v_muoValPixHits)[i];
			(*inputVec)[19] = (double)(*v_muoNTrkVHits)[i];
			(*inputVec)[20] = (double)(*v_muoGNchi2)[i];
			(*inputVec)[21] = (double)(*v_muoVMuHits)[i];
			(*inputVec)[22] = (double)(*v_muoNumMatches)[i];
			(*inputVec)[23] = (double)(*v_muoPFiso)[i];
			(*inputVec)[24] = (double)(*v_muoQprod)[i];

			// retrive the classifier responses
			if(TMath::Abs((*v_muoLund)[i])==13) histSgn->Fill(classReader->GetMvaValue( *inputVec ));
			if(TMath::Abs((*v_muoLund)[i])!=13) histBkg->Fill(classReader->GetMvaValue( *inputVec ));
			
		}
	}

	cout << "=== Macro        : Event loop done! " << endl;

	//DRAWING

	TCanvas *c1 = new TCanvas();
	TCanvas *c2 = new TCanvas();

	int nSgn = histSgn->GetEntries();
	int nBkg = histBkg->GetEntries();
	int lower_bin = 0;
	int higher_bin = histBins;

	while(histBkg->GetBinContent(lower_bin)==0){
		lower_bin++;
	}

	while(histSgn->GetBinContent(higher_bin)==0){
		higher_bin--;
	}

	int offset = 5;

	if(higher_bin+=offset<histBins){
		higher_bin+=offset;
	}else{
		higher_bin=histBins;
	}

	if(lower_bin-=offset>0){
		lower_bin-=offset;
	}else{
		lower_bin=0;
	}

	int n = 50+offset*2;
	float pass=(histSgn->GetBinCenter(higher_bin)-histSgn->GetBinCenter(lower_bin))/float(n); 
	float MvaCut = 0;

	float *rocX = new float[n]; //sgn eff
	float *rocY = new float[n]; //bkg eff


	for(int i=0; i<n; ++i){
		float MvaCut=histSgn->GetBinCenter(lower_bin)+pass*i;
		float nSgn_ = 0;
		float nBkg_ = 0;
		int bin = lower_bin;

		while(histSgn->GetBinCenter(bin)<MvaCut){
			nSgn_+=histSgn->GetBinContent(bin);
			nBkg_+=histBkg->GetBinContent(bin);
			bin+=1;
		}

		rocX[i]=(nSgn-nSgn_)/nSgn;
		rocY[i]=(nBkg_)/nBkg;

	}

	c1->cd();
	histSgn->SetLineColor(kBlue);
	histBkg->SetLineColor(kRed);
	histSgn->SetLineWidth(2);
	histBkg->SetLineWidth(2);

	histSgn->Scale(1/(double)nSgn);
	histBkg->Scale(1/(double)nBkg);

	histSgn->GetXaxis()->SetRangeUser(histSgn->GetBinCenter(lower_bin), histSgn->GetBinCenter(higher_bin));

	histSgn->Draw("HIST");
	histBkg->Draw("HIST SAME");

	c2->cd();
	TGraph *roc = new TGraph(n, rocX, rocY);
	roc->Draw("AL");
	roc->SetTitle("ROC Curve");
	roc->GetXaxis()->SetRangeUser(0.0, 1.0);
	roc->GetXaxis()->SetTitle("Signal Efficiency");
	roc->GetYaxis()->SetTitle("Background Rejection");
	roc->SetMaximum(1.0);
	roc->SetMinimum(0.0);
	roc->SetLineWidth(2);
	roc->Draw("AL");
	c2->SetGridx();
	c2->SetGridy();

}