import ROOT as r
import root_numpy as rn
import pickle
import numpy as np
import matplotlib.pyplot as plt
from xgboost import XGBClassifier
from tools.data_processing import *
from tools.mc_fitter import MC_fit
import os
import argparse

l_flv=['e','mu']
mother_ID=["Ds","Dplus","both"]



def select_and_fit(mother_index_fit, l_index, x_cut, iteration, test):

	#PATHS NEEDED
	
	FIT_PATH=l_flv[l_index]+'/fits'
	BDT_PATH=l_flv[l_index]+'/BDTs/test_'+str(test)

	with open(BDT_PATH+'/variables_used.pickle', 'rb') as f:  
	        branches_needed_0 = pickle.load(f)
	branches_needed=['cos_thetal']+branches_needed_0

	#Data loading
	MC_Dplus_sig_dict, MC_Ds_sig_dict, _ = load_datasets(l_index)

	with open('/disk/lhcb_data/davide/Rphipi/NN_for_selection/'+l_flv[l_index]+l_flv[l_index]+'/'+'data_for_NN_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f:
	    data_dict=pickle.load(f)
	with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[0]+'_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f:
	    MC_Ds_mass_for_fit=pickle.load(f)
	with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[1]+'_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f:
	    MC_Dplus_mass_for_fit=pickle.load(f)


	"""
	Normalising the chi2s in data and MC

	"""
	MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict = norm_chi2(MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict)


	"""

	Applying mass cuts on data, 
	retrieving mass distributions for sig MC and data


	"""


	data_mass, mc_mass, data_dict, lower_cut, upper_cut =mass_cut_for_fit(mother_index_fit=mother_index_fit, l_index=l_index,
													 					  branches_needed=branches_needed, 
													 					  data_dict=data_dict, MC_Dplus_sig_dict=MC_Dplus_sig_dict, MC_Ds_sig_dict=MC_Ds_sig_dict)

	"""

	Load the MC DCB shape parameters, 
	if not present fit them

	"""

	if os.path.exists(FIT_PATH+'/MC_'+mother_ID[0]+'_CB_params.pickle'):
		with open(FIT_PATH+'/MC_'+mother_ID[0]+'_CB_params.pickle','rb') as f:
			
			Ds_CB_params = pickle.load(f)

			alpha_Ds_l=Ds_CB_params['alpha_Ds_l']
			power_Ds_l=Ds_CB_params['n_Ds_l']
			alpha_Ds_r=Ds_CB_params['alpha_Ds_r']
			power_Ds_r=Ds_CB_params['n_Ds_r']
			frac_Ds=Ds_CB_params['CB fraction']

	else:

		MC_Ds_fit_results, alpha_Ds_l, power_Ds_l, alpha_Ds_r, power_Ds_r, frac_Ds = MC_fit( MC_Ds_mass_for_fit, mother_index=0, l_index=l_index)

	if os.path.exists(FIT_PATH+'/MC_'+mother_ID[1]+'_CB_params.pickle'):
		with open(FIT_PATH+'/MC_'+mother_ID[1]+'_CB_params.pickle','rb') as f:
			
			Dplus_CB_params = pickle.load(f)

			alpha_Dplus_l=Dplus_CB_params['alpha_Dplus_l']
			power_Dplus_l=Dplus_CB_params['n_Dplus_l']
			alpha_Dplus_r=Dplus_CB_params['alpha_Dplus_r']
			power_Dplus_r=Dplus_CB_params['n_Dplus_r']
			frac_Dplus=Dplus_CB_params['CB fraction']

	else:

		MC_Dplus_fit_results, alpha_Dplus_l, power_Dplus_l, alpha_Dplus_r, power_Dplus_r, frac_Dplus = MC_fit(MC_Dplus_mass_for_fit, mother_index=1, l_index=l_index)

	"""

	Preprocess data for XGBoost classifier


	"""

	data_2, mc_2, data_mean, mc_mean, data_std, mc_std = preprocess_for_XGBoost(MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict, branches_needed, mother_index_fit=mother_index_fit)


	Nsig_from_MC=mc_mass.shape[0]


	"""

	BDT selection


	"""

	print("BDT cut at {0} fit".format(x_cut))

	print("BDT selection..")

	k = np.random.randint(10)
	MODEL_PATH=BDT_PATH+'/XG_'+str(k)

	loaded_model = pickle.load(open(MODEL_PATH+"/XG_"+str(k)+"_.pickle.dat", "rb"))

	output_XG=loaded_model.predict_proba(data_2)
	output_XG_signal=loaded_model.predict_proba(mc_2)

	data_selected=np.array(data_mass[np.where(output_XG[:,1]>x_cut)])
	XG_mc_selected=np.array(mc_mass[np.where(output_XG_signal[:,1]>x_cut)])

	sig_sel_eff=np.float(XG_mc_selected.shape[0])/np.float(Nsig_from_MC)

	print("Signal selection efficiency: {0}".format(sig_sel_eff))

	"""

	Select mass window for fit


	"""


	cut_indices=[]
	for i in range(len(data_selected)):
		if lower_cut<data_selected[i]<upper_cut:

			cut_indices.append(i)
    
	data_cut=data_selected[cut_indices]/1000.   

	lower_cut/=1000.
	upper_cut/=1000.

	"""

	Create a rootfile from which will do the unbinned fit


	"""

	data_array=np.array(

			[data_cut[i] for i in range(len(data_cut))], dtype=[('D_reco_M', np.float32)]

			)




	rn.array2root(data_array,
					filename='/disk/lhcb_data/davide/Rphipi/selected_data/'+l_flv[l_index]+l_flv[l_index]+'/sel_data_'+l_flv[l_index]+l_flv[l_index]+'.root',
					treename='decay_tree',
					mode='recreate',
					)
	f=r.TFile('/disk/lhcb_data/davide/Rphipi/selected_data/'+l_flv[l_index]+l_flv[l_index]+'/sel_data_'+l_flv[l_index]+l_flv[l_index]+'.root')
	tree=f.Get("decay_tree") 

	#some tricks to pass Roofit a dataset in python
	mass = np.array([0],dtype=np.float32)

	branch = tree.GetBranch("D_reco_M")
	branch.SetAddress(mass)

	num_entries=tree.GetEntries()
	m = r.RooRealVar("D mass reco (GeV/c^2)","D mass reco (GeV/c^2)",lower_cut,upper_cut)
	l = r.RooArgSet(m)
	data_set = r.RooDataSet("data set", "data set", l)

	for i in range(num_entries):
		tree.GetEvent(i)
		r.RooAbsRealLValue.__assign__(m, mass[0])
		data_set.add(l, 1.0)
	f.Close()
	"""

	Retrieve the needed pdfs


	"""

	##Creating D+ signal PDF left tail
	mean_Dplus= r.RooRealVar("mean_Dplus","mean_Dplus",1.87,1.83,1.91)
	sigma_Dplus = r.RooRealVar("sigma_Dplus","sigma_Dplus",0.020,0.,0.2)

	sig_frc_Dplus =r.RooRealVar("Dplus Double crystalball fraction","Dplus Double crystalball fraction",frac_Dplus)

	al_Dplus_left = r.RooRealVar("alpha_Dplus_left","alpha_Dplus_left",alpha_Dplus_l)
	pwr_Dplus_left = r.RooRealVar("n_Dplus_left","n_Dplus_left",power_Dplus_l)
	sig_Dplus_left = r.RooCBShape("signal Dplus left","signal Dplus left", m, mean_Dplus, sigma_Dplus, al_Dplus_left, pwr_Dplus_left)
	#Creating D+ signal PDF right tail

	al_Dplus_right = r.RooRealVar("alpha_Dplus_right","alpha_Dplus_right",alpha_Dplus_r)
	pwr_Dplus_right = r.RooRealVar("n_Dplus_right","n_Dplus_right",power_Dplus_r)
	sig_Dplus_right = r.RooCBShape("signal Dplus right","signal Dplus right", m, mean_Dplus, sigma_Dplus, al_Dplus_right, pwr_Dplus_right)

	#Adding the two CBs with their relative fraction
	sig_Dplus = r.RooAddPdf("signal Dplus","Dplus mass peak",r.RooArgList(sig_Dplus_left, sig_Dplus_right),r.RooArgList(sig_frc_Dplus))

	#Creating Ds+ signal PDF
	mean_Ds= r.RooRealVar("mean_Ds","mean_Ds",1.97,1.93,2.01)
	sigma_Ds = r.RooRealVar("sigma_Ds","sigma_Ds",0.020,0.,0.2)

	sig_frc_Ds =r.RooRealVar("Ds Double crystalball fraction","Ds Double crystalball fraction",frac_Ds)

	al_Ds_left = r.RooRealVar("alpha_Ds_left","alpha_Ds_left",alpha_Ds_l)
	pwr_Ds_left = r.RooRealVar("n_Ds_left","n_Ds_left",power_Ds_l)
	sig_Ds_left = r.RooCBShape("signal Ds left","signal Ds left", m, mean_Ds, sigma_Ds, al_Ds_left, pwr_Ds_left)
	#Creating D+ signal PDF right tail

	al_Ds_right = r.RooRealVar("alpha_Ds_right","alpha_Ds_right",alpha_Ds_r)
	pwr_Ds_right = r.RooRealVar("n_Ds_right","n_Ds_right",power_Ds_r)
	sig_Ds_right = r.RooCBShape("signal Ds right","signal Ds right", m, mean_Ds, sigma_Ds, al_Ds_right, pwr_Ds_right)

	#Adding the two CBs with their relative fraction
	sig_Ds = r.RooAddPdf("sig_Ds","Ds mass peak",r.RooArgList(sig_Ds_left, sig_Ds_right),r.RooArgList(sig_frc_Ds))

	#Model the background

	coef0 = r.RooRealVar("c0","coefficient #0",1.0,-10.,10)
	coef1 = r.RooRealVar("c1","coefficient #1",0.1,-10.,10)
	coef2 = r.RooRealVar("c2","coefficient #2",-0.1,-10.,10)
	bkg = r.RooChebychev("bkg","background p.d.f.",m,r.RooArgList(coef0,coef1,coef2))

	#add it altogether

    #Add 2 sources of signal
	if mother_ID[mother_index_fit]=='both':

		NDs = r.RooRealVar("nDs","nDs",100.,0.,40000.)
		NDplus = r.RooRealVar("nDplus","nDplus",100.,0.,40000.)
		nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.)
		model = r.RooAddPdf("model","Dplus and Ds mass peaks",r.RooArgList(sig_Ds, sig_Dplus, bkg),r.RooArgList(NDs, NDplus,nbkg))


	if mother_ID[mother_index_fit]=='Ds':
		NDs = r.RooRealVar("nDs","nDs",100.,0.,40000.)
		nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.)
		model = r.RooAddPdf("model","Ds mass peak",r.RooArgList(sig_Ds, bkg),r.RooArgList(NDs,nbkg))

	if mother_ID[mother_index_fit]=='Dplus':

		NDplus = r.RooRealVar("nDplus","nDplus",100.,0.,40000.)
		nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.)
		model = r.RooAddPdf("model","Dplus mass peak",r.RooArgList(sig_Dplus, bkg),r.RooArgList(NDplus,nbkg))

	fitr = model.fitTo(data_set,r.RooFit.Extended(),r.RooFit.Save())
	xframe = m.frame(r.RooFit.Title("Fit to "+l_flv[l_index]+" data"))

	bkg_component = r.RooArgSet(bkg)
	sig_Dplus_component = r.RooArgSet(sig_Dplus)
	sig_Ds_component = r.RooArgSet(sig_Ds)
	#sig_D_component = r.RooArgSet(sig_Dplus_component,sig_Ds_component)

	data_set.plotOn(xframe)
	model.plotOn(xframe)
	hpull=xframe.pullHist()

	n_param = fitr.floatParsFinal().getSize()
	reduced_chi_square = xframe.chiSquare(n_param)

	model.plotOn(xframe,r.RooFit.Components(bkg_component),r.RooFit.LineStyle(2))

	if mother_ID[mother_index_fit]=='Dplus':

		model.plotOn(xframe, r.RooFit.Components(sig_Dplus_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) )

	if mother_ID[mother_index_fit]=='Ds':

		model.plotOn(xframe, r.RooFit.Components(sig_Ds_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) )

	if mother_ID[mother_index_fit]=='both':

		model.plotOn(xframe, r.RooFit.Components(sig_Ds_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) )
		model.plotOn(xframe, r.RooFit.Components(sig_Dplus_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) )


	model.paramOn(xframe, r.RooFit.Layout(0.69,0.99,0.92), r.RooFit.Format("NEU", r.RooFit.AutoPrecision(1)))
	xframe.getAttText().SetTextSize(0.04)

	xframe2 = m.frame(r.RooFit.Title("Pulls"))
	xframe2.addPlotable(hpull,"P")

	c = r.TCanvas("Fit {0}".format(iteration),"Fit {0}".format(iteration),900,600)
	pad1 = r.TPad("pad1", "pad1", 0, 0.35, 1, 1.0)
	pad1.SetBottomMargin(0)
	pad1.Draw()
	c.cd()
	pad2 = r.TPad("pad2", "pad2", 0, 0.05, 1, 0.35)
	pad2.SetTopMargin(0)
	pad2.SetBottomMargin(0.2)
	pad2.SetGridx()
	pad2.SetGridy()
	pad2.Draw()

	pad1.cd()
	xframe.Draw()
	#xframe.GetYaxis().SetLabelOffset(-0.01)
	pad2.cd()
	xframe2.Draw()
	xframe2.GetXaxis().SetTitleSize(0.09)
	xframe2.GetXaxis().SetLabelSize(0.1)
	xframe2.GetYaxis().SetLabelSize(0.05)



	print("chi2 {0}".format(reduced_chi_square))
	c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}/fit_{1}.png".format(iteration,iteration))
	
	
	xframe_sw = m.frame(r.RooFit.Title("Weighted "+l_flv[l_index]+" data"))

	if mother_ID[mother_index_fit]=='both':

		splot = r.RooStats.SPlot(
		'sData','sData', data_set, model,
		r.RooArgList(NDs, NDplus, nbkg)

		)

		Ds_sw=splot.GetSWeightVars()[0]
		Dplus_sw=splot.GetSWeightVars()[1]
		bkg_sw=splot.GetSWeightVars()[2]

		argset = r.RooArgSet(m, Ds_sw)
		dataset_Ds_w = r.RooDataSet("Ds weighted dataset","Ds weighted dataset",data_set, argset,"","nDs_sw")

		argset = r.RooArgSet(m, Dplus_sw)
		dataset_Dplus_w = r.RooDataSet("Dplus weighted dataset","Dplus weighted dataset",data_set, argset,"","nDplus_sw")

		
		dataset_Ds_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kBlue))
		dataset_Dplus_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kRed))


		Ds_sw.setRange(-1.2,1.5)
		frame_Ds_sw = Ds_sw.frame(r.RooFit.Title("signal Ds sWeights"))
		data_set.plotOn(frame_Ds_sw, r.RooFit.MarkerSize(0.05))

		Dplus_sw.setRange(-1.2,1.5)
		frame_Dplus_sw = Dplus_sw.frame(r.RooFit.Title("signal Dplus sWeights"))
		data_set.plotOn(frame_Dplus_sw, r.RooFit.MarkerSize(0.05))
		   
		bkg_sw.setRange(-.6,1.2)
		frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights"))
		data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05))
		   

		sframe = m.frame(r.RooFit.Title("m vs sWeights"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(Ds_sw),
		                r.RooFit.MarkerColor(r.kGreen), r.RooFit.Name("Ds_sig"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(Dplus_sw),
		                r.RooFit.MarkerColor(r.kBlue), r.RooFit.Name("Dplus_sig"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw),
		                r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg"))

		legend = r.TLegend(0.89, 0.89, 0.5, 0.7)
		legend.AddEntry(sframe.findObject('Ds_sig'), 'Signal Ds sWeights', 'p')
		legend.AddEntry(sframe.findObject('Dplus_sig'), 'Signal Dplus sWeights', 'p')
		legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p')
		#
		c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600)
		xframe_sw.Draw()
		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration))


		c = r.TCanvas("sWeights {0}".format(iteration),"sWeights {0}".format(iteration),900,900)
		c.Divide(2,2)
		c.cd(1)
		frame_Ds_sw.Draw()
		c.cd(2)
		frame_Dplus_sw.Draw()
		c.cd(3)
		frame_bkg_sw.Draw()
		c.cd(4)
		sframe.Draw()
		legend.Draw()
		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration))


		nbkg_w = splot.GetYieldFromSWeight("nbkg_sw")
		nDs_w = splot.GetYieldFromSWeight("nDs_sw")
		nDplus_w = splot.GetYieldFromSWeight("nDplus_sw")
		raw_sig=nDs_w+nDplus_w

	if mother_ID[mother_index_fit]=='Ds':

		splot = r.RooStats.SPlot(
		'sData','sData', data_set, model,
		r.RooArgList(NDs, nbkg)

		)

		Ds_sw=splot.GetSWeightVars()[0]
		bkg_sw=splot.GetSWeightVars()[1]

		argset = r.RooArgSet(m, Ds_sw)
		dataset_Ds_w = r.RooDataSet("Ds weighted dataset","Ds weighted dataset",data_set, argset,"","nDs_sw")

		dataset_Ds_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kBlue))

		Ds_sw.setRange(-3.,1.4)
		frame_Ds_sw = Ds_sw.frame(r.RooFit.Title("signal Ds sWeights"))
		data_set.plotOn(frame_Ds_sw, r.RooFit.MarkerSize(0.05))
		   
		bkg_sw.setRange(-.5,5.)
		frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights"))
		data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05))
		   
		sframe = m.frame(r.RooFit.Title("m vs sWeights"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(Ds_sw),
		                r.RooFit.MarkerColor(r.kGreen), r.RooFit.Name("Ds_sig"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw),
		                r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg"))

		legend = r.TLegend(0.89, 0.89, 0.75, 0.8)
		legend.AddEntry(sframe.findObject('Ds_sig'), 'Signal Ds sWeights', 'p')
		legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p')
		#
		c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600)
		xframe_sw.Draw()
		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration))


		c = r.TCanvas("sWeights {0}".format(iteration),"sWeights  {0}".format(iteration),900,600)
		c.cd()
		pad1 = r.TPad("pad1", "pad1", 0, 0.5, 0.5, 1.0)
		pad1.Draw()
		c.cd()
		pad2 = r.TPad("pad2", "pad2", 0.5, 0.5, 1.0, 1.0)
		pad2.Draw()
		c.cd()
		pad3 = r.TPad("pad3", "pad3", 0, 0.05, 1, 0.50)
		pad3.SetGridx()
		pad3.Draw()
		

		pad1.cd()
		frame_Ds_sw.Draw()
		pad2.cd()
		frame_bkg_sw.Draw()
		pad3.cd()
		sframe.Draw()
		legend.Draw()

		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration))


		nbkg_w = splot.GetYieldFromSWeight("nbkg_sw")
		nDs_w = splot.GetYieldFromSWeight("nDs_sw")
		raw_sig=nDs_w

	if mother_ID[mother_index_fit]=='Dplus':

		splot = r.RooStats.SPlot(
		'sData','sData', data_set, model,
		r.RooArgList(NDplus, nbkg)

		)

		Dplus_sw=splot.GetSWeightVars()[0]
		bkg_sw=splot.GetSWeightVars()[1]

		argset = r.RooArgSet(m, Dplus_sw)
		dataset_Dplus_w = r.RooDataSet("Dplus weighted dataset","Dplus weighted dataset",data_set, argset,"","nDplus_sw")

		#model.plotOn(xframe_sw, r.RooFit.Components(sig_Dplus_component),r.RooFit.LineColor(2),r.RooFit.LineStyle(2) )
		dataset_Dplus_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kRed))


		Dplus_sw.setRange(-1.2,1.8)
		frame_Dplus_sw = Dplus_sw.frame(r.RooFit.Title("signal Dplus sWeights"))
		data_set.plotOn(frame_Dplus_sw, r.RooFit.MarkerSize(0.05))
		   
		bkg_sw.setRange(-1.2,2.1)
		frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights"))
		data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05))
		   

		sframe = m.frame(r.RooFit.Title("m vs sWeights"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(Dplus_sw),
		                r.RooFit.MarkerColor(r.kBlue), r.RooFit.Name("Dplus_sig"))
		data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw),
		                r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg"))

		legend = r.TLegend(0.89, 0.89, 0.75, 0.8)
		legend.AddEntry(sframe.findObject('Dplus_sig'), 'Signal Dplus sWeights', 'p')
		legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p')
		#
		c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600)
		xframe_sw.Draw()
		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration))


		c = r.TCanvas("sWeights {0}".format(iteration),"sWeights  {0}".format(iteration),900,600)
		c.cd()
		pad1 = r.TPad("pad1", "pad1", 0, 0.5, 0.5, 1.0)
		pad1.Draw()
		c.cd()
		pad2 = r.TPad("pad2", "pad2", 0.5, 0.5, 1.0, 1.0)
		pad2.Draw()
		c.cd()
		pad3 = r.TPad("pad3", "pad3", 0, 0.05, 1, 0.50)
		pad3.SetGridx()
		pad3.Draw()
		

		pad1.cd()
		frame_Dplus_sw.Draw()
		pad2.cd()
		frame_bkg_sw.Draw()
		pad3.cd()
		sframe.Draw()
		legend.Draw()

		c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration))


		nbkg_w = splot.GetYieldFromSWeight("nbkg_sw")
		nDplus_w = splot.GetYieldFromSWeight("nDplus_sw")
		raw_sig=nDplus_w

	fom=(Nsig_from_MC)/np.sqrt(nbkg_w+Nsig_from_MC)

	eff_corrected_sig = (raw_sig)/(sig_sel_eff)

	fit_results={
				 '# XGBoost test {0} BDT output'.format(test): '\n',
				 '# background': nbkg_w,
				 #'err bkg': err_bkg,
				 '# fitted signal': raw_sig,
				 #'error sig': raw_sig_err,
				 '# of signal in MC': Nsig_from_MC,
				 'signal selection efficiency': sig_sel_eff,
				 'S/sqrt(S+B)': fom,
				 'efficiency corrected yeld': eff_corrected_sig,
				 'chi 2': reduced_chi_square
				 }

	return fit_results


if __name__=='__main__':

	
	parser = argparse.ArgumentParser(description='Fit at different BDT cuts')
	parser.add_argument('-iteration', metavar='Fit iteration', type =int, help='Fit iteration')
	#parser.add_argument('-test', metavar='BDT test', type =int, help='Which BDT test you are using')
	parser.add_argument('-x_cut', metavar='BDT cut', type =float, help='BDT cut')
	parser.add_argument('-l_index', metavar='Lepton flavour', type =int, help='Which lepton flavour are you fitting?')
	parser.add_argument('-mother_index_fit', metavar='Decay mother', type =int, help='Which mass window are you fitting?')
	args = parser.parse_args()
	
	test=4

	#test=args.test
	i = args.iteration
	l_index = args.l_index
	mother_index_fit = args.mother_index_fit
	x_cut=args.x_cut

	
	FIT_PATH=l_flv[l_index]+'/fits'
	BDT_PATH=l_flv[l_index]+'/BDTs/test_'+str(test)


	if not os.path.exists(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/{0}/'.format(i)):
		os.mkdir(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/{0}/'.format(i))

	
	fit_results = select_and_fit(mother_index_fit, l_index, x_cut, i, test)

	with open(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/fit_results_{0}.txt'.format(i), 'w') as g:
		print('Saving fit results...')
		for key, value in fit_results.items():
			g.write('%s:%s\n' % (key, value))
	g.close()