#include <algorithm>
#include "TFile.h"
#include "TTree.h"
#include "TLeaf.h"
#include "TString.h"
#include "TH1F.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TEventList.h"
#include "TDirectory.h"
#include <iostream>
#include "TLine.h"
#include "TLatex.h"
#include "TStyle.h"
#include "TObjArray.h"
#include "TPRegexp.h"
#include "TRegexp.h"
#include <map>
#include <utility>
#include <algorithm>

#define MAX 1000
#define BINS 500
#define NCAN 3
#define MAXVARS 8  //maximum allowed number of variables for one selection
#define MAXCUTS 12  //maximum allowed number of cuts for one selection

typedef std::vector<pair<double, int> > ValueMap;
typedef std::map<int, int> CountMap;

CountMap evcnt, evcntrec[10];
std::map<TString, TString> mctvar;

TString vars[MAX];
TString cuts[MAX];
double cut[MAX];
double qual[MAX];
int idx[MAX];

double N0_sig, N0_bg;
double Nsigev, Nbgev;
bool dstarmode;

Float_t fbranch[MAX];
Int_t   ibranch[MAX];
Bool_t  bbranch[MAX];
Int_t ev, run, mode, rec, nbranch;

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

int gettype(TTree *t, TString varname)
{
	TString leaftype = t->GetLeaf(varname)->GetTypeName();
	
	if (leaftype=="Float_t")	 return 0; 
	else if (leaftype=="Int_t")	 return 1;  
	else if (leaftype=="Bool_t") return 2;
	
	return -1;
}

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

bool mycompare(int i, int j)
{
  return qual[i]>qual[j];
}

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

int init(TTree *t)
{
	t->SetBranchAddress("ev",&ev);
	t->SetBranchAddress("run",&run);
	t->SetBranchAddress("mode",&mode);
	t->SetBranchAddress("recmode",&rec);
	
	TObjArray* branches = t->GetListOfBranches();	
	
	nbranch = 0;
	
	for(int i=0; i<=branches->GetLast(); ++i)
	{
		TBranch* branch = (TBranch*)branches->UncheckedAt(i);
		TString v=branch->GetName(); 

		if ( v=="ev" || v=="mode" || v=="run" || v=="nsig" ) continue;
		if ( v.Contains("pdg") || v.Contains("beam") || v.Contains("mct") ) continue;
		if ( v.BeginsWith("t") && v!="thr") continue;
		if ( v.EndsWith("vx") || v.EndsWith("vy") || v.EndsWith("vz") || v.EndsWith("pocmag")) continue;

		bool ok=false;
		
		if ( v.BeginsWith("es") || v=="mmiss" || v.EndsWith("d0m") ) ok=true;
		if ( v.EndsWith("p") || v.EndsWith("tht") || v.EndsWith("pcm") || v.EndsWith("thtcm") || v.EndsWith("pt") ) ok=true;
		if ( ok || v.EndsWith("ang") || v.Contains("poc") ) ok=true;
		if ( ok || v.Contains("pid") || v.Contains("min") || v.Contains("max") || v.Contains("sum") || v.Contains("fw") ) ok=true;
		if ( ok || v.EndsWith("sph") || v.EndsWith("apl") || v.EndsWith("pla") || v.EndsWith("thr") || v.EndsWith("cir") ) ok=true;
		
		if ( v=="xmdif" && dstarmode ) ok = true;
		if ( v=="xd0m" && dstarmode ) ok = false;
		
		if (!ok) continue;
		
		Int_t dtype=gettype(t, v);
		switch (dtype)
		{
			case 0: t->SetBranchAddress(v, &(fbranch[nbranch])); break;
			case 1: t->SetBranchAddress(v, &(ibranch[nbranch])); break;
			case 2: t->SetBranchAddress(v, &(bbranch[nbranch])); break;
		}
		
		vars[nbranch] = v;
		nbranch++;
	}
	
	return nbranch;
}

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

int uid(int lev, int lrun, int lmode)
{
	return lev+10000*lrun+(((lmode/100)%10)*20+lmode%10)*100000;
}

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

int countEvents(TTree *t, TEventList &el)
{
	t->SetEventList(0);
	t->SetBranchStatus("*",0);
	t->SetBranchStatus("ev",1);
	t->SetBranchStatus("run",1);
	t->SetBranchStatus("mode",1);
	t->SetBranchStatus("recmode",1);
		
	evcnt.clear();
	for (int j=0;j<10;++j) evcntrec[j].clear();
	
	for (int i=0;i<el.GetN();++i)
	{
		t->GetEntry(el.GetEntry(i));
		evcnt[uid(ev,run,mode)]+=1;
		if (rec<10) evcntrec[rec][uid(ev,run,mode)]+=1;
	}
	t->SetBranchStatus("*",1);

	return evcnt.size();
}
// ---------------------------------------------------------------

void makeMaps(TTree *t, TString varname, TEventList &els, TEventList &elb, ValueMap &sig, ValueMap &bg, int id)
{
	int i;
	t->SetBranchStatus("*",0);
	t->SetBranchStatus("ev",1);
	t->SetBranchStatus("mode",1);
	t->SetBranchStatus("run",1);
	t->SetBranchStatus(varname,1);
	
	Float_t var;
	Int_t dtype=gettype(t, varname);
	
	sig.clear();
	bg.clear();

	// prepare the pairs of variable value, eventnumber for signal and count signals
	for (i=0;i<els.GetN();++i)
	{
		t->GetEntry(els.GetEntry(i));
		switch (dtype) {case 2: var = (Float_t)bbranch[id]; break; case 1: var = (Float_t) ibranch[id]; break; default: var=fbranch[id]; }
		
		sig.push_back(std::make_pair(var, uid(ev,run,mode)  ));
	}
		
	// prepare the pairs of variable value, eventnumber for background and count background
	for (i=0;i<elb.GetN();++i)
	{
		t->GetEntry(elb.GetEntry(i));
		switch (dtype) {case 2: var = (Float_t)bbranch[id]; break; case 1: var = (Float_t) ibranch[id]; break; default: var=fbranch[id]; }
		
		bg.push_back(std::make_pair(var, uid(ev,run,mode) ));
	}
}

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

double bestEffEvt(TTree *t, TString varname, TEventList &els, TEventList &elb, double &bestcut, double supr, int id)
{	
	CountMap bgcnt, sigcnt, sigcnt2;	
	ValueMap sigvals, bgvals;
	
	int Nbgval = elb.GetN();
	int dtype = gettype(t,varname);
	
	makeMaps(t, varname, els, elb, sigvals, bgvals, id);
	
	
	// sort signals by first element = variable value
	sort(bgvals.begin(), bgvals.end());

	int i=0;
	bgcnt.clear();
	// find cut for current efficiency requirement left cut
	while ( (bgcnt.size()/(double)Nbgev)<(1-supr) && i<Nbgval ) bgcnt[bgvals[i++].second]+=1;
	double leftcut = bgvals[i].first;
	
	bgcnt.clear();
	i=bgvals.size()-1;
	// find cut for current efficiency requirement right cut
	while ( (bgcnt.size()/(double)Nbgev)<(1-supr) && i>=0) bgcnt[bgvals[i--].second]+=1; 
	double rightcut = bgvals[i].first;
	
	sigcnt.clear();
	sigcnt2.clear();
	
	for (i=0;i<(int)sigvals.size();++i)
	{
		if (sigvals[i].first<=leftcut)  sigcnt[sigvals[i].second]+=1;
		if (sigvals[i].first>=rightcut) sigcnt2[sigvals[i].second]+=1;
	}
	
	t->SetBranchStatus("*",1);
		
	double lefteff  = sigcnt.size()/Nsigev;
	double righteff = sigcnt2.size()/Nsigev;

	//cout <<varname<<" "<<flush;
	
	TString prec="%.4f";
	
	bestcut = rightcut;
	if (lefteff>righteff) 
	{
		bestcut = leftcut;
		
		if (bestcut<1e-3 || fabs(1.-bestcut)<1e-3) prec="%.6f";
		if (dtype==0) cuts[id] = TString::Format(TString("%s<="+prec).Data(),varname.Data(),leftcut);
		else cuts[id] = TString::Format("%s<=%.1f",varname.Data(),leftcut);
		
		qual[id] = lefteff;

		return lefteff;
	}
	
	if (bestcut<1e-3 || fabs(1.-bestcut)<1e-3) prec="%.6f";

	if (dtype==0) cuts[id] = TString::Format(TString("%s>="+prec).Data(),varname.Data(),rightcut);
	else cuts[id] = TString::Format("%s>=%.1f",varname.Data(),rightcut);

	qual[id] = righteff;

	return righteff;
}

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

double bestSuppressionEvt(TTree *t, TString varname, TEventList &els, TEventList &elb, double &bestcut, double eff, int id)
{	
	CountMap bgcnt, bgcnt2, sigcnt;	
	ValueMap sigvals, bgvals;
	
	int Nsigval = els.GetN();
	int dtype = gettype(t,varname);
	
	makeMaps(t, varname, els, elb, sigvals, bgvals, id);
	
	// sort signals by first element = variable value
	sort(sigvals.begin(), sigvals.end());

	int i=0;
	sigcnt.clear();
	// find cut for current efficiency requirement left cut
	while ( (sigcnt.size()/(double)Nsigev)<eff && i<Nsigval ) sigcnt[sigvals[i++].second]+=1;
	double leftcut = sigvals[i].first;
	
	sigcnt.clear();
	i=sigvals.size()-1;
	// find cut for current efficiency requirement right cut
	while ( (sigcnt.size()/(double)Nsigev)<eff && i>=0) sigcnt[sigvals[i--].second]+=1; 
	double rightcut = sigvals[i].first;
	
	bgcnt.clear();
	bgcnt2.clear();
	
	for (i=0;i<(int)bgvals.size();++i)
	{
		if (bgvals[i].first<=leftcut)  bgcnt[bgvals[i].second];
		if (bgvals[i].first>=rightcut) bgcnt2[bgvals[i].second];
	}
	
	t->SetBranchStatus("*",1);
	
	
	double leftsupr = (Nbgev-bgcnt.size())/Nbgev;
	double rightsupr = (Nbgev-bgcnt2.size())/Nbgev;

	//cout <<varname<<" "<<flush;
	
	if (leftcut!=leftcut || rightcut!=rightcut) return 0;
	
	TString prec="%.4f";
	
	bestcut = rightcut;
	if (leftsupr>rightsupr) 
	{
		bestcut = leftcut;
		
		if (bestcut<1e-3 || fabs(1.-bestcut)<1e-3) prec="%.6f";
		if (dtype==0) cuts[id] = TString::Format(TString("%s<="+prec).Data(),varname.Data(),leftcut);
		else cuts[id] = TString::Format("%s<=%.1f",varname.Data(),leftcut);
		
		qual[id] = leftsupr;

		return leftsupr;
	}
	
	if (bestcut<1e-3 || fabs(1.-bestcut)<1e-3) prec="%.6f";

	if (dtype==0) cuts[id] = TString::Format(TString("%s>="+prec).Data(),varname.Data(),rightcut);
	else cuts[id] = TString::Format("%s>=%.1f",varname.Data(),rightcut);

	qual[id] = rightsupr;

	return rightsupr;
}

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

int findcut(TTree *t, TEventList &els, TEventList &elb, double supr, double &bestqa)
{
	int i;
	
	bestqa = -999.;
	
	TString bestcut="";
	
	int bestid = -1;
	
	//cout <<"000/000"<<flush;

	for(i=0; i<nbranch; ++i)
	{
		idx[i] = i;
		
		if (i%(nbranch/10)==0){ cout <<"#"<<flush;}
//		printf("\b\b\b\b\b\b\b%03d/%03d",i+1,nbranch);
//		fflush(stdout);
		
		double qa=0;
		
		if (supr<0)	qa = bestEffEvt(t, vars[i], els, elb, cut[i], -supr, i);
		else        qa = bestSuppressionEvt(t, vars[i], els, elb, cut[i], supr, i);	
		
		if (qa>bestqa && fabs(cut[i])>1e-5 && fabs(cut[i]-1.0)>1e-5) 
		{
			bestqa  = qa;
			bestcut = cuts[i];
			bestid = i;
		}
	}
	
	return bestid;
}

// ---------------------------------------------------------------
int countVars(TString s)
{
	int cnt = 0;
	TString svar="";
	TRegexp rvar("-[a-zA-Z_][a-zA-Z0-9_]+-");
	
	TObjArray *tok = s.Tokenize("&&");
	int N = tok->GetEntries();	
	for (int i=0;i<N;++i) 
	{
		if (i<100) 
		{
			TString toks = ((TObjString*)tok->At(i))->String();
			toks.ReplaceAll("\t","");
			toks = toks.Strip(TString::kBoth);
			if (toks.Index("<")>0) toks = toks(0,toks.Index("<"));
			if (toks.Index(">")>0) toks = toks(0,toks.Index(">"));
			svar += "-"+toks+"-";
		}
	}
	while (svar(rvar)!="") {svar.ReplaceAll(svar(rvar),""); cnt++;}
	return cnt;
}

// ---------------------------------------------------------------
int countCuts(TString s)
{
	TObjArray *tok = s.Tokenize("&&");
	return tok->GetEntries();	
}

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

int autocutx(TString fname, TString precut="", double supr=0.95, double target=0.0001, double mineff = 0.1, double minreleff = 0.0, int evmult=10000, double norm=1.0, int n0s=-1)
{
// 	gStyle->SetTitleX(0.2);
// 	gStyle->SetTitleY(0.993);
// 	gStyle->SetTitleH(0.07);
// 	gStyle->SetPadTopMargin(0.08);
// 	gStyle->SetTitleBorderSize(0);
	gStyle->SetOptStat(0);
	
	TCanvas *c1=new TCanvas("c1","c1",10,10,1800,550);
	c1->Divide(7,2);
	
	TString tagcut = "tag";
	
	TRegexp ren("M[0-9][0-9][0-9]"); 
	TRegexp regntp("n[0-9]+");
	TRegexp regS("[0-9]+S");
	TRegexp regB("[0-9]+B");
	
	TString en = fname(ren);  en = en(1,5); 
	TString s  = fname(regS); s  = s(0,s.Length()-1);
	TString b  = fname(regB); b  = b(0,b.Length()-1);
	
	TString ntp   = fname(regntp);
	TString smode = ntp(1,ntp.Length());
	int dmode     = smode.Atoi()/10;
	dstarmode     = (dmode==11 || dmode==13 || dmode==15);
	
	if (dmode>=10 && dmode<20) norm = 0.5;
	
	if (dmode/10 == 6) norm = 0.2;
	
	if (n0s<0) n0s = s.Atoi();
	int n0b = b.Atoi();
	
	N0_sig = n0s*evmult;
	N0_bg = n0b*evmult;
	
	cout <<"tree '"<<ntp.Data()<<"' with N_sig = "<<N0_sig<<", N_bg = "<<N0_bg<<endl;

	TFile *f=new TFile(fname,"READ");
	TTree *t=(TTree*)f->Get(ntp);
	
	init(t);
	
	TEventList elsall("elsall"), elball("elball");
	TEventList els("els"), elb("elb");
	
	if (precut=="") precut = tagcut;
	else precut = tagcut+"&&"+precut;
	
	TString sigcut = "xmct";
	TString bgcut  = "!xmct";
	
	t->Draw(">>elsall",tagcut+"&&"+sigcut);
	t->Draw(">>elball",tagcut+"&&"+bgcut);
	float nsig = countEvents(t,elsall);
	float nbg = countEvents(t,elball);
	
	double beff=1., seff=1., rseff=1.;
	
	TRegexp rnum("[\\-]*[0-9]+\\.[0-9]+$");
	int cnt=0;
	
	while (cnt<150 && beff>target && seff>=mineff && rseff>=minreleff) 
	{		
		c1->cd(cnt%7+1); 
		t->SetLineColor(1);	t->Draw("xm",precut);
		t->SetLineColor(2);	t->Draw("xm",precut+"&&"+sigcut,"same");
		t->SetLineColor(4);	t->Draw("xm",precut+"&&"+bgcut,"same");
		
		c1->Update();
				
		t->Draw(">>els",sigcut+"&&"+precut);
		t->Draw(">>elb",bgcut+"&&"+precut);

		Nbgev  = countEvents(t, elb);
		Nsigev = countEvents(t, els);
		
	   	beff  = Nbgev/N0_bg;
	   	seff  = Nsigev/N0_sig;
		rseff = Nsigev/nsig;
		
		printf("S=%6d  B=%6d (%.2f): ",(int)Nsigev, (int)Nbgev, supr);
		
		// stop if target reached or signal efficiency too small
		if (beff<=target || seff<mineff || rseff<minreleff) continue;
		
		double qa=0.;
		
		int bestid = findcut(t, els, elb, supr, qa);
		
		TString bestcut = cuts[bestid];
		
		std::vector<int> myidx (idx, idx+nbranch);	
		std::sort(myidx.begin(), myidx.end(), mycompare);	
		
		if (precut.Length()>3 && precut.Index(vars[bestid]+">")<0 && precut.Index(vars[bestid]+"<")<0)
		{
			int i = 0, ncuts = countCuts(precut);
			
			while (i<nbranch && precut.Index(vars[idx[i]]+">")<0 && precut.Index(vars[idx[i]]+"<")<0) ++i;
			
			if (qual[idx[i]]/qa>0.8 ) bestcut = cuts[idx[i]];
		}		
		
		TString thenum = bestcut(rnum);
		double cutval = thenum.Atof();
		TString thevar = bestcut(0,bestcut.Index(thenum));
		TString thebarevar = thevar;
		thebarevar.ReplaceAll(">","");
		thebarevar.ReplaceAll("<","");
		thebarevar.ReplaceAll("=","");
				
		TRegexp rcut(thevar+"[\\-]*[0-9]+\\.[0-9]+");
		TString theoldcut = precut(rcut);
		
		c1->cd(cnt%7+8);
		
		t->SetLineColor(1);	t->Draw(thebarevar,precut);
		t->SetLineColor(2);	t->Draw(thebarevar,precut+"&&"+sigcut,"same");
		t->SetLineColor(4);	t->Draw(thebarevar,precut+"&&"+bgcut,"same");
		c1->Update();
		TLine l; l.SetLineColor(6); l.SetLineWidth(2);
		l.DrawLine(cutval,0,cutval,gPad->GetUymax()*0.9);
		
		c1->Update();
		
		if (theoldcut!="") 
			precut.ReplaceAll(theoldcut,bestcut);
		else
			precut += "&&"+bestcut;		
			
		cout <<" : "<<precut<<endl;
		cnt++;
	}
	
	cout <<"\nCUT : "<<en<<smode<<" : "<<precut.Data()<<endl;
	cout <<"SIG EVT: "<<Nsigev<<" ev    BG: "<<Nbgev<<" ev"<<endl;

	printf("SIG EFF : %6.1f%%    BG EFF : %7.3f%%\n",Nsigev/N0_sig*100.,Nbgev/N0_bg*100.);
	printf("SIG REL : %6.1f%%    BG REL : %7.3f%%\n\n",Nsigev/nsig*100.,Nbgev/nbg*100.);

	cout<<"Recoil : "; for (int k=0;k<10;++k) printf("    %02d ",k);cout <<endl;
	cout<<"Eff    : "; for (int k=0;k<10;++k) printf("%6.1f%%",(double) evcntrec[k].size()/evmult*norm*100.);cout <<endl<<endl;

	t->SetEventList(0);
  return 0;
}