import ROOT as r
import root_numpy as rn
import pickle
import numpy as np
import matplotlib.pyplot as plt
import array as array
from xgboost import XGBClassifier
from tools.data_processing import *

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

def MC_fit(mc_mass, mother_index, l_index):
    FIT_PATH=l_flv[l_index]+'/fits'
    print("Fitting MC "+mother_ID[mother_index]+" signal")


    if l_flv[l_index]=='mu':
        if mother_ID[mother_index]=='Dplus':
            lower_mass_cut=mc_mass.min()
            upper_mass_cut=mc_mass.max()
        if mother_ID[mother_index]=='Ds':
            lower_mass_cut=mc_mass.min()
            upper_mass_cut=mc_mass.max()
    
    if l_flv[l_index]=='e':
        if mother_ID[mother_index]=='Dplus':
            lower_mass_cut=mc_mass.min()
            upper_mass_cut=mc_mass.max()
        if mother_ID[mother_index]=='Ds':
            lower_mass_cut=mc_mass.min()
            upper_mass_cut=mc_mass.max()   

    mc_fit=mc_mass/1000.
    lower_mass_cut/=1000.
    upper_mass_cut/=1000.
    
    #create a rootfile from which will do the unbinned fit



    data_array=np.array([mc_fit[i] for i in range(len(mc_fit))], dtype=[('MC_D_m', np.float32)])
    
    rn.array2root(data_array,
                  filename='/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[mother_index]+'_'+l_flv[l_index]+l_flv[l_index]+'_mass.root',
                  treename='D_M',
                  mode='recreate',
                 )
    f=r.TFile('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[mother_index]+'_'+l_flv[l_index]+l_flv[l_index]+'_mass.root')
    tree=f.Get("D_M") 
    
    #some tricks to pass Roofit a dataset in python
    mass = np.array([0],dtype=np.float32)
    branch = tree.GetBranch("MC_D_m")
    branch.SetAddress(mass)
    
    num_entries=tree.GetEntries()
    m = r.RooRealVar("m","m",lower_mass_cut,upper_mass_cut)
    l = r.RooArgSet(m)
    data_hist = r.RooDataSet("data histogram", "data", l)
    
    
    for i in range(num_entries):
        tree.GetEvent(i)
        r.RooAbsRealLValue.__assign__(m, mass[0])
        data_hist.add(l, 1.0)
    
    
    #Creating the signal PDF as two crystall balls
    
    if l_flv[l_index]=='mu':
        if mother_ID[mother_index]=='Dplus':
            mean_d_startval=1.85
            mean_d_up=1.95
            mean_d_down=1.75
        if mother_ID[mother_index]=='Ds':
            mean_d_startval=1.95
            mean_d_up=2.05
            mean_d_down=1.85
        
    if l_flv[l_index]=='e':
        if mother_ID[mother_index]=='Dplus':
            mean_d_startval=1.85
            mean_d_up=2.30
            mean_d_down=1.70
        if mother_ID[mother_index]=='Ds':
            mean_d_startval=1.95
            mean_d_up=2.30
            mean_d_down=1.70

    
    #Creating D signal PDF left tail
    mean_d= r.RooRealVar("mean_D","mean_D",mean_d_startval,mean_d_down,mean_d_up)
    sigma_d = r.RooRealVar("sigma_D","sigma_D",0.020,0.,0.2)
    
    alpha_l = r.RooRealVar("alpha_D_l","alpha_D_l",-0.020,-10.,10.)
    n_l = r.RooRealVar("n_D_l","n_D_l",0.020,-50.,50.)
    sig_l = r.RooCBShape("signal_D_l","signal_D_l", m, mean_d, sigma_d, alpha_l, n_l)

    #Creating D signal PDF right tail
  
    alpha_r = r.RooRealVar("alpha_D_r","alpha_D_r",0.020,-10.,10.)
    n_r = r.RooRealVar("n_D_r","n_D_r",0.020,-50.,50.)
    sig_r = r.RooCBShape("signal_D_r","signal_D_r", m, mean_d, sigma_d, alpha_r, n_r)

    
    #Add 2 sources of signal
    # first add the sigs together 
    frac = r.RooRealVar("# of Ds sig events","# of Ds sig events",0.5,0.,1.)
    
    model = r.RooAddPdf("model","two mass peaks",r.RooArgList(sig_l,sig_r),r.RooArgList(frac))
    
    fitr = model.fitTo(data_hist,r.RooFit.Save())

    xframe = m.frame(r.RooFit.Title("Fit to "+l_flv[l_index]+" MC signal"))

    data_hist.plotOn(xframe)
    model.plotOn(xframe)
    c = r.TCanvas()
    xframe.Draw()
    
    n_param = fitr.floatParsFinal().getSize()
    reduced_chi_square = xframe.chiSquare(n_param)
    
    CB_fraction=frac.getVal()
    Nsig_from_MC=mc_mass.shape[0]
    
    a_l = alpha_l.getVal()
    pwr_l= n_l.getVal()
    
    a_r= alpha_r.getVal()
    pwr_r=n_r.getVal()
    
    print("chi2 {0}".format(reduced_chi_square))
    c.SaveAs(l_flv[l_index]+"/fits/MC_"+mother_ID[mother_index]+"_fit.png")
    
    MC_fit_results={
                 '# Fitting the MC signal':'\n',
                 'CB fraction': CB_fraction,
                 '# of signal in MC': Nsig_from_MC,
                 'alpha_'+mother_ID[mother_index]+'_l': a_l, 
                 'n_'+mother_ID[mother_index]+'_l': pwr_l, 
                 'alpha_'+mother_ID[mother_index]+'_r': a_r, 
                 'n_'+mother_ID[mother_index]+'_r': pwr_r,
                 }
    print('Saving MC fit results...')
    with open(FIT_PATH+'/MC_'+mother_ID[mother_index]+'_fit_results.txt', 'w') as f:
            
            for key, value in MC_fit_results.items():
                f.write('%s:%s\n' % (key, value))
    with open(FIT_PATH+'/MC_'+mother_ID[mother_index]+'_CB_params.pickle','w') as handle:
            
            pickle.dump(MC_fit_results, handle)
    print('Saved MC fit results \n Proceeding to fit script')
    return MC_fit_results, a_l, pwr_l, a_r, pwr_r, CB_fraction