import ROOT
#from ROOT import TTree, TFile, Double
import numpy as np
from pdg_const import pdg
import Tkinter
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import pickle as pkl
import sys
import time
from helperfunctions import display_time, prepare_plot
import cmath as c

mmu = pdg['muon_M']
mb = pdg["bquark_M"]
ms = pdg["squark_M"]
mK = pdg["Ks_M"]
mB = pdg["Bplus_M"]

class model:

    def __init__(self):

        #Initialize constants

        self.mmu = pdg['muon_M']
        self.mb = pdg["bquark_M"]
        self.ms = pdg["squark_M"]
        self.mK = pdg["Ks_M"]
        self.mB = pdg["Bplus_M"]

        self.C7eff = pdg["C7eff"]
        self.C9eff = pdg["C9eff"]
        self.C10eff = pdg["C10eff"]

        #self.C1 = pdg["C1"]
        #self.C2 = pdg["C2"]
        #self.C3 = pdg["C3"]
        #self.C4 = pdg["C4"]

        self.GF = pdg["GF"] #Fermi coupling const.
        self.alpha_ew = pdg["alpha_ew"]
        self.Vts = pdg["Vts"]
        self.Vtb = pdg["Vtb"]

        self.param_str = []
        self.param_val = []

        #helperconstants

        self.x_min = 2*self.mmu
        self.x_max = (self.mB - self.mK) - 0.1
        self.total_pdf_string = "self.total_nonres(q2)"
        self.mode = ""
        self.param_str.append("total_scale_amp")
        self.param_val.append(1.0)
        self.param_str.append("_mean")
        self.param_val.append(0.0)
        self.res_counter = 0
        self.reco_steps = 100000

        #Cusp info

        self.param_str.append("cusp mean")
        self.param_val.append(1.0)
        self.param_str.append("cusp sigma_L")
        self.param_val.append(1.0)
        self.param_str.append("cusp sigma_R")
        self.param_val.append(1.0)
        self.param_str.append("cusp amp")
        self.param_val.append(0.0)

    def formfactor(self, q2, subscript): #returns real value

        #check if subscript is viable

        if subscript != "0" and subscript != "+" and subscript != "T":
            raise ValueError('Wrong subscript entered, choose either 0, + or T')

        #get constants

        mh = self.mK
        mbstar0 = pdg["mbstar0"]
        mbstar = pdg["mbstar"]
        b0 = pdg["b0"]
        bplus = pdg["bplus"]
        bT = pdg["bT"]

        #N comes from derivation in paper

        N = 3

        #some helperfunctions

        tpos = (self.mB - self.mK)**2
        tzero = (self.mB + self.mK)*(np.sqrt(self.mB)-np.sqrt(self.mK))**2

        z_oben = np.sqrt(tpos - q2) - np.sqrt(tpos - tzero)
        z_unten = np.sqrt(tpos - q2) + np.sqrt(tpos - tzero)
        z = z_oben/z_unten

        #calculate f0

        if subscript == "0":
            prefactor = 1/(1 - q2/(mbstar0**2))
            _sum = 0

            for i in range(N):
                _sum += b0[i]*(z**i)

            return prefactor * _sum

        #calculate f+ or fT

        else:
            prefactor = 1/(1 - q2/(mbstar**2))
            _sum = 0

            if subscript == "T":
                b = bT
            else:
                b = bplus

            for i in range(N):
                _sum += b[i] * (z**i - ((-1)**(i-N)) * (i/N) * z**N)

            return prefactor * _sum

    def axiv_nonres(self, q2): #returns real value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #Some helperfunctions

        beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2))

        kabs = np.sqrt(mB**2 + q2**2/mB**2 + mK**4/mB**2 - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2)

        #prefactor in front of whole bracket

        prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)

        #left term in bracket

        bracket_left = 2./3. * kabs**2 * beta**2 * abs(C10eff*self.formfactor(q2, "+"))**2

        #middle term in bracket

        _top = 4. * mmu**2 * (mB**2 - mK**2) * (mB**2 - mK**2)

        _under = q2 * mB**2

        bracket_middle = _top/_under * abs(C10eff * self.formfactor(q2, "0"))**2

        #Note sqrt(q2) comes from derivation as we use q2 and plot q

        return self.param_val[0]*prefactor1 * (bracket_left + bracket_middle) * 2 * np.sqrt(q2)

    def vec_nonres(self, q2): #returns real value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C7eff = self.C7eff
        C9eff = self.C9eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #Some helperfunctions

        beta = np.sqrt(abs(1. - 4. * self.mmu**2. / q2))

        kabs = np.sqrt(mB**2 + q2**2/mB**2 + mK**4/mB**2 - 2 * (mB**2 * mK**2 + mK**2 * q2 + mB**2 * q2) / mB**2)

        #prefactor in front of whole bracket

        prefactor1 = GF**2. *alpha_ew**2. * (abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)

        #right term in bracket

        prefactor2 = kabs**2 * (1. - 1./3. * beta**2)

        abs_bracket = abs(C9eff * self.formfactor(q2, "+") + 2 * C7eff * (mb + ms)/(mB + mK) * self.formfactor(q2, "T"))**2

        bracket_right = prefactor2 * abs_bracket

        #Note sqrt(q2) comes from derivation as we use q2 and plot q

        return self.param_val[0]*prefactor1 * bracket_right * 2 * np.sqrt(q2)

    def total_nonres(self, q2): #returns complex value

        #Load constants

        GF = self.GF
        alpha_ew = self.alpha_ew
        Vtb = self.Vtb
        Vts = self.Vts
        C10eff = self.C10eff
        C9eff = self.C9eff
        C7eff = self.C7eff

        mmu = self.mmu
        mb = self.mb
        ms = self.ms
        mK = self.mK
        mB = self.mB

        #vector nonresonant part

        vec_nonres_part = self.vec_nonres(q2)

        #axial verctor nonresonant part including C7

        axiv_nonres_part = self.axiv_nonres(q2)

        #Complete term

        return complex(vec_nonres_part + axiv_nonres_part, 0.0)

    def generate_points(self, set_size = 10000, x_min = 2* mmu, x_max = (mB - mK) - 0.1, save = True, verbose = 1, mode = "true_data", resolution = 7.0, nonres_set_size = 44000): #returns part_set (=dic)

        #Check if entered mode and verbose are valid

        if mode != "slim_points" and mode != "fast_binned" and mode != "true_data":
            raise ValueError('Invalid mode entered, choose either slim_points, fast_binned or true_data')

        if verbose != 0 and verbose != 1:
            raise ValueError('Invalid verbose entered, choose either 0 (no verbose) or 1 (verbose)')

        #Setup contants and variables

        self.mode = mode

        mB = self.mB
        mK = self.mK
        mmu = self.mmu

        #Range of function in MeV

        dq = np.linspace(x_min, x_max ,self.reco_steps)

        #Translate to MeV**2

        dgamma_tot = []

        #Prepare some variables used in all modes

        for i in dq:
            dgamma_tot.append(self.total_pdf(i**2))

        dq2 = []

        for i in dq:
                dq2.append(i**2)

        _max = max(dgamma_tot)

        #slim_mode: Generates data of a set size in range(x_min,x_max) according to total_pdf -> Only saves the valid q values (Hit and miss method)

        if mode == "slim_points":

            #Prepare variables

            x_part = []

            print("Generating set of size {}...".format(int(set_size)))

            #Change seed here if necessary

            # ROOT.TRandom1().SetSeed(0)

            #Preparecalls for verbose percentage counter

            if verbose != 0:
                verbose_calls = []
                j = 0
                o = 0
                for j in range(100):
                    verbose_calls.append(int(set_size*(j+1)/100))

            #Get start time for projected time and other variables

            start = time.time()
            counter = 0
            counter_x = 0

            #Loop until set has wanted size

            while len(x_part) < set_size:
                counter += 1
                x = ROOT.TRandom1().Uniform(x_min, x_max)
                y = ROOT.TRandom1().Uniform(0, _max)

                if y < self.total_pdf(x**2):
                    x_part.append(x)
                    counter_x += 1

                #progress calls
                if verbose != 0:
                    end = time.time()
                    if o < 100 and counter%5000 == 0:
                        print(" {0}{1} completed".format(o, "%"))
                        print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(counter_x+1)-(end-start)))))
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                    if o >=100:
                        print(" Time to generate set: {0}".format(display_time(int(end-start))))

                    if len(x_part) == verbose_calls[o]:
                        o += 1

                #Generate random points

                counter += 1
                x = ROOT.TRandom1().Uniform(x_min, x_max)
                y = ROOT.TRandom1().Uniform(0, _max)

                #Add particles if under pdf

                if y < self.total_pdf(x**2):
                    x_part.append(x)
                    counter_x += 1

                #progress calls

                if verbose != 0:
                    end = time.time()
                    if o < 100 and counter%5000 == 0:
                        print(" {0}{1} completed".format(o, "%"))
                        print(" Projected time left: {0}".format(display_time(int((end-start)*set_size/(counter_x+1)-(end-start)))))
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                    if o >=100:
                        print(" Time to generate set: {0}".format(display_time(int(end-start))))

                    if counter_x == verbose_calls[o]:
                        o += 1

            print(" {0} out of {1} particles chosen!".format(int(counter_x), counter))

            print(" Set generated!")

            #Save the set

            if save:

                part_set = {"x_part" : x_part}

                pkl.dump( part_set, open("./data/set_{0}_range({1}-{2}).pkl".format(int(set_size),int(x_min), int(x_max)) , "wb" ) )

                print(" Set saved!")

            print

            #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)

            return part_set

        elif mode == "true_data":

            #Prepare variables

            x_part = []

            print("Generating set with {0} total rare nonresonant particles...".format(int(nonres_set_size)))

            #Change Seed here if necessary

            #ROOT.TRandom1().SetSeed(0)

            #Range of the whole nonres spectrum (used to count total hits in the nonres part)

            x_min_nr = self.x_min
            x_max_nr = self.x_max

            #Prepare verbose calls

            if verbose != 0:
                verbose_calls = []
                j = 0
                o = 0
                while j < 100:
                    j += verbose
                    verbose_calls.append(int(nonres_set_size*j/100))

            #Take starting time and other variables for projected time

            start = time.time()

            counter = 0
            counter_x = 0
            counter_x_nr = 0

            #Loop until enough particles in nonresonant part

            while counter_x_nr < nonres_set_size:

                #Generate random points

                counter += 1
                x = ROOT.TRandom1().Uniform(x_min_nr, x_max_nr)
                y = ROOT.TRandom1().Uniform(0, _max)

                #Add the part to the set if within range(x_min,x_max) and under pdf

                if x > x_min and x < x_max and y < self.total_pdf(x**2):
                    x_part.append(x)
                    counter_x += 1

                #Count the hits in the nonres part -> Range of the whole pdf!

                if y < abs(self.total_nonres(x**2)):
                    counter_x_nr += 1

                #progress calls
                if verbose != 0:
                    end = time.time()
                    if o < 100 and counter%5000 == 0:
                        print(" {0}{1} completed".format(o, "%"))
                        print(" Projected time left: {0}".format(display_time(int((end-start)*nonres_set_size/(counter_x_nr+1)-(end-start)))))
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                        sys.stdout.write("\033[F")
                        sys.stdout.write("\x1b[2K")
                    if o >=100:
                        print(" Time to generate set: {0}".format(display_time(int(end-start))))

                    if counter_x_nr == verbose_calls[o]:
                        o += 1


            print(" {0} out of {1} particles chosen!".format(int(counter_x), counter))

            print(" Set generated!")

            #Save the set

            if save:

                part_set = {"x_part" : x_part, "counter_x_nr": counter_x_nr}

                pkl.dump( part_set, open("./data/set_true_tot_nonres_{0}_range({1}-{2}).pkl".format(int(nonres_set_size),int(x_min), int(x_max)) , "wb" ) )

                print(" Set saved!")

            print

            #returns all the chosen points (x_part, y_part) and all the random points generated (x, y)

            return part_set
        elif mode == "fast_binned":

            #Calculate number of bins

            nbins = int((x_max-x_min)/resolution)

            print("Generating set with {} bins...".format(nbins))

            #Get the mean dq values for each bin and get value of pdf there

            dq = np.linspace(x_min, x_max ,nbins+1)

            bin_mean = []
            bin_true_height = []

            for i in range(len(dq)-1):
                a = dq[i]
                b = dq[i+1]
                c = (a+b)/2
                bin_mean.append(c)

                height = self.total_pdf(c**2)
                bin_true_height.append(height)

            #Scale true bin size according to nonres_set_size

            dq_scan = np.linspace(self.x_min, self.x_max, self.reco_steps)

            dgamma_scan = []

            for i in dq_scan:
                dgamma_scan.append(abs(self.total_nonres(i**2)))

            _mean_scan = np.mean(dgamma_scan)

            _mean_histo_nonres = nonres_set_size/nbins

            _ =  _mean_histo_nonres/_mean_scan

            # print(_)

            for i in range(len(bin_true_height)):
                bin_true_height[i] = bin_true_height[i]*_

            #Start time

            start = time.time()

            bin_height = []

            #Draw random numbers

            for i in range(len(bin_mean)):
                bin_height.append(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))
                #print(int(ROOT.TRandom1().Gaus(bin_true_height[i], np.sqrt(bin_true_height[i]))))

            #progress calls
            if verbose != 0:
                end = time.time()
                print(" Time to generate set: {0}s".format(end-start))

            print(" {0} bins simulated".format(nbins))

            print(" Set generated!")

            #Save the set

            if save:

                _ = 0

                for i in bin_height:
                    _ += i

                part_set = {"bin_mean" : bin_mean, "bin_height": bin_height, "nbins": nbins, "bin_true_height": bin_true_height}

                pkl.dump( part_set, open("./data/binned_set_{0}bins_{1}part.pkl".format(nbins, _) , "wb" ) )

                print(" Set saved!")

            print

            return part_set

    def draw_plots(self, part_set, mode, x_min = 2* mmu, x_max = (mB - mK) - 0.1, resolution = 7.0): #draws and saves plots, returns nothing
        #Resolution based on experiment chosen to be ~7MeV
        #check for invalid entries

        if mode != "slim_points" and mode != "fast_binned" and mode != "true_data":
            raise ValueError('Wrong mode entered, choose either slim_points, fast_binned or true_data' )
        if mode != self.mode:
            raise ValueError('self.mode and mode are different, set them to the same value')

        print("Generating plots")

        if mode == "fast_binned":

            #Load variables

            mB = self.mB
            mK = self.mK
            mmu = self.mmu

            #Range of function in MeV

            #Choosing the number of steps greatly varies the mean of dgamma!!! 5000 -> 10000 yields a factor of 1.68
            dq = np.linspace(x_min, x_max ,self.reco_steps)

            #Translate to MeV**2

            dq2 = []

            for i in dq:
                dq2.append(i**2)

            #calculate formfactors

            ff_plus = []
            ff_T = []
            ff_0 = []

            for i in dq:
                ff_0.append(self.formfactor(i**2, "0"))
                ff_T.append(self.formfactor(i**2, "T"))
                ff_plus.append(self.formfactor(i**2, "+"))

            #calculate nonresonant

            dgamma_axiv_nonres = []
            dgamma_vec_nonres = []
            dgamma_total_nonres = []
            dgamma_tot = []

            for i in dq:
                    dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
                    dgamma_vec_nonres.append(self.vec_nonres(i**2))
                    dgamma_total_nonres.append(abs(self.total_nonres(i**2)))
                    dgamma_tot.append(self.total_pdf(i**2))


            #Plot formfactors

            plt.clf()

            plt.plot(dq2, ff_0, label = "0")
            plt.plot(dq2, ff_T, label = "T")
            plt.plot(dq2, ff_plus, label = "+")

            prepare_plot("Formfactors")

            plt.savefig("./plots/fast_binned/ff.png")

            print(" ff.png created")

            #Plot nonresonant part

            plt.clf()

            plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
            plt.plot(dq, dgamma_vec_nonres, label = "vec")
            plt.plot(dq, dgamma_total_nonres, label = "total_nonres")

            prepare_plot("Nonresonant parts")

            plt.savefig("./plots/fast_binned/vec_axiv.png")

            print(" vec_axiv.png created")

            #Plot all pdfs and parts

            tot_y = []
            jpsi_y = []
            psi2s_y = []
            total_nonres_y = []
            cusp_y = []

            jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]

            psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]

            for i in range(len(dq)):
                #print(i**2 - 4*(mmu**2))
                tot_y.append(abs(self.total_pdf(dq[i]**2)))
                jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)))
                psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)))
                total_nonres_y.append(abs(self.total_nonres(dq[i]**2)))
                cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.param_val[2], self.param_val[5], self.param_val[3], self.param_val[4] )))
                #resonance(self, q2, _mass, width, phase, scale):
                #w[i] = np.sqrt(w[i])
                #print(test_y[i])

            plt.clf()

            plt.plot(dq, tot_y, label = "total pdf")
            plt.plot(dq, jpsi_y, label = "jpsi")
            plt.plot(dq, psi2s_y, label = "psi2s")
            plt.plot(dq, total_nonres_y, label = "nonres")
            plt.plot(dq, cusp_y, label = "cusp")

            prepare_plot("All pdfs")

            plt.ylim(0, 2*self.param_val[1])

            plt.savefig("./plots/fast_binned/pdf_and_parts.png")

            print(" pdf_and_parts.png created")

            #Create histo with pdf

            dq2 = []

            for i in dq:
                    dq2.append(i**2)

            dgamma_tot = []

            for i in dq2:
                dgamma_tot.append(self.total_pdf(i))

            #Load particle set

            bin_mean = part_set["bin_mean"]
            bin_height = part_set["bin_height"]

            #Scale the bins to the pdf

            _sum = np.sum(bin_height)

            nbins = part_set["nbins"]

            _bin_mean = np.mean(bin_height)
            pdf_mean = self.param_val[1]

            for i in range(len(bin_height)):
                bin_height[i] = bin_height[i] * pdf_mean/_bin_mean

            plt.clf()

            plt.ylim(0.0,self.param_val[1]*2)

            #Choose range in between the 2 resonances
            # plt.xlim(3150, 3650)

            #Draw the plot

            plt.hist(bin_mean, bins=nbins, range=(x_min, x_max), weights = bin_height, label = "toy data binned")

            plt.plot(dq, dgamma_tot, label = "pdf")

            prepare_plot("{0} random points generated ({1} particles)".format(len(bin_mean), int(_sum)))

            plt.savefig("./plots/fast_binned/histo.png")

            print(" histo.png created")

            print(" All plots drawn \n")

            return

        else:

            #Range of function in MeV

            dq = np.linspace(x_min, x_max ,self.reco_steps)

            #Translate to MeV**2

            dq2 = []

            for i in dq:
                    dq2.append(i**2)

            #calculate formfactors

            ff_plus = []
            ff_T = []
            ff_0 = []

            for i in dq:
                ff_0.append(self.formfactor(i**2, "0"))
                ff_T.append(self.formfactor(i**2, "T"))
                ff_plus.append(self.formfactor(i**2, "+"))

            #calculate nonresonant

            dgamma_axiv_nonres = []
            dgamma_vec_nonres = []
            dgamma_total_nonres = []
            dgamma_tot = []

            for i in dq:
                    dgamma_axiv_nonres.append(self.axiv_nonres(i**2))
                    dgamma_vec_nonres.append(self.vec_nonres(i**2))
                    dgamma_total_nonres.append(abs(self.total_nonres(i**2)))
                    dgamma_tot.append(self.total_pdf(i**2))


            #Plot formfactors

            plt.clf()

            plt.plot(dq2, ff_0, label = "0")
            plt.plot(dq2, ff_T, label = "T")
            plt.plot(dq2, ff_plus, label = "+")

            prepare_plot("Formfactors")

            plt.savefig("./plots/points/ff.png")

            print(" ff.png created")


            #Plot nonresonant part

            plt.clf()

            plt.plot(dq, dgamma_axiv_nonres, label = "axiv")
            plt.plot(dq, dgamma_vec_nonres, label = "vec")
            plt.plot(dq, dgamma_total_nonres, label = "total_nonres")

            prepare_plot("Nonresonant parts")

            plt.savefig("./plots/points/vec_axiv.png")

            print(" vec_axiv.png created")

            #Particle set

            x_part = part_set["x_part"]

            set_size = len(x_part)

            #Histo unnormalized

            bins = int((x_max-x_min)/resolution)

            plt.clf()

            _y, _x, _ = plt.hist(x_part, bins=bins, range=(x_min, x_max), label = "toy data binned ({0} points)".format(set_size))

            _mean_histo = float(np.mean(_y))

            prepare_plot("Binned toy data")

            plt.savefig("./plots/points/histo_raw.png")

            print(" histo_raw.png created")


            #Histo and pdf normailzed

            plt.clf()

            #Plot histo

            #Scale the histo to the pdf

            _ = []

            for i in range(len(_y)):
                _y[i] = _y[i]*np.mean(dgamma_tot)/_mean_histo
                _.append(x_min+(x_max-x_min)/bins*i)


            plt.hist(_, bins=bins, range=(x_min, x_max), weights = _y, label = "toy data binned  (scaled down to pdf)")

            plt.plot(dq, dgamma_tot, label = "pdf")

            #Only show from 0 to twice the mean of pdf
            # plt.ylim(0, 2*self.param_val[1])

            prepare_plot("{0} random points generated according to pdf".format(set_size))

            plt.savefig("./plots/points/histo.png")

            print(" histo.png created")

            #Plot all pdfs

            tot_y = []
            jpsi_y = []
            psi2s_y = []
            total_nonres_y = []
            cusp_y = []


            jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]

            psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]

            for i in range(len(dq)):
                #print(i**2 - 4*(mmu**2))
                tot_y.append(abs(self.total_pdf(dq[i]**2)))
                jpsi_y.append(abs(self.resonance(dq[i]**2, jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)))
                psi2s_y.append(abs(self.resonance(dq[i]**2, psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)))
                total_nonres_y.append(abs(self.total_nonres(dq[i]**2)))
                cusp_y.append(abs(self.bifur_gauss(dq[i]**2, self.param_val[2], self.param_val[5], self.param_val[3], self.param_val[4] )))
                #resonance(self, q2, _mass, width, phase, scale):
                #w[i] = np.sqrt(w[i])
                #print(test_y[i])

            plt.clf()

            plt.plot(dq, tot_y, label = "total pdf")
            plt.plot(dq, jpsi_y, label = "jpsi")
            plt.plot(dq, psi2s_y, label = "psi2s")
            plt.plot(dq, total_nonres_y, label = "nonres")
            plt.plot(dq, cusp_y, label = "cusp")

            prepare_plot("All pdfs and parts")

            plt.savefig("./plots/points/pdf_and_parts.png")

            print(" pdf_and_parts.png created")

            print(" All plots drawn \n")

            return

    def total_pdf(self, q2): #return absolut value of sum of all complex values (real number)

        #Calculate the pdf with the added resonances

        exec("_sum = abs({0})".format(self.total_pdf_string))

        return _sum

    def resonance(self, q2, _mass, width, phase, scale): #returns complex value

        #Helper variables

        p = 0.5 * np.sqrt(q2 - 4*(mmu**2))

        p0 =  0.5 * np.sqrt(_mass**2 - 4*mmu**2)

        gamma_j = p / p0 * _mass /q2 * width

        #Calculate the resonance

        _top = complex(_mass * width, 0.0)

        _bottom = complex((_mass**2 - q2), -_mass*gamma_j)

        com = _top/_bottom * scale

        #Rotate by the phase

        r = abs(com)

        _phase = c.phase(com)

        _phase += phase

        x = c.cos(phase)*r
        y = c.sin(phase)*r

        com = complex(x,y)

        return self.param_val[0]*com

    def add_resonance(self, _mass, width, phase, scale, name): #adds string to total_pdf and adds constant places for fit

        #Count the number of resonances added to the pdf -> used to keep track of the parameters

        _ = self.res_counter

        #Adds the resonace to the pdf in form of a string (to be executed later)

        self.param_str.append(name + " mass")
        self.param_val.append(_mass)
        self.param_str.append(name + " width")
        self.param_val.append(width)
        self.param_str.append(name + " phase")
        self.param_val.append(phase)
        self.param_str.append(name + " scale")
        self.param_val.append(scale)

        self.total_pdf_string += "+ self.resonance(q2,{0},{1},{2},{3})".format(self.param_val[6+_*4], self.param_val[7+_*4],self.param_val[8+_*4], self.param_val[9+_*4])

        self.res_counter += 1

    def bifur_gauss(self, q2, mean, amp, sigma_L, sigma_R): #returns complex value

        #Get q out of q2

        q = np.sqrt(q2)

        #Choose the right sigma depending on side of the cusp

        if q < mean:
            sigma = sigma_L
        else:
            sigma = sigma_R

        #Calculate the exponential part of the cusp

        _exp = np.exp(- (q-mean)**2 / (2 * sigma**2))

        #Scale so the total area under curve is 1 and the top of the cusp is continuous

        dgamma = amp*_exp/(np.sqrt(2*np.pi))*2*(sigma_L*sigma_R)/(sigma_L+sigma_R)

        com = complex(dgamma, 0)

        return self.param_val[0]*com

    def add_cusp(self, mean, amp, sigma_L, sigma_R): #adds string to total_pdf and adds constant places for fit

        #Add cusp to total_pdf

        self.total_pdf_string += "+ self.bifur_gauss(q2,{0},{1},{2},{3})".format(mean, "self.param_val[5]", "self.param_val[3]", "self.param_val[4]")

        #Save the variables of the cusp -> Used for fitting

        self.param_val[2] = mean
        self.param_val[3] = sigma_L
        self.param_val[4] = sigma_R
        self.param_val[5] = amp

    def normalize_pdf(self, step_scan = False): #Normalizes pdf with a global factor in front, saves mean and global factor

        print("Normalizing pdf")

        #step_scan = True: Scans for the optimal number of steps to scan over the function to scale the whole pdf (done numerically)

        if step_scan:

            print(" Scanning for optimal step size to scan pdf...")

            #Scan over different number of steps

            mean_list = []

            for i in range(11):
                steps = 50000 + i*10000
                x_scan = np.linspace(self.x_min, self.x_max, steps)

                y_scan = []

                for i in x_scan:
                    y_scan.append(self.total_pdf(i**2))

                mean_list.append(np.mean(y_scan))

            _mean = np.mean(mean_list)

            #Choose the one closest to the mean of all step variations

            diff = 1.0
            _ = 0
            for i in range(11):
                if abs(mean_list[i]-_mean) < diff:
                    reco_steps = steps = 50000 + i*10000
                    _ = i

            #Save recommended step size for later

            self.reco_steps = reco_steps

            print("  Optimal number of steps to scan pdf: {0}".format(reco_steps))

            _mean = mean_list[_]

        else:

            #Scan over pdf to normalize with reco_steps

            x_scan = np.linspace(self.x_min, self.x_max, self.reco_steps)

            y_scan = []

            for i in x_scan:
                y_scan.append(self.total_pdf(i**2))

            _mean = np.mean(y_scan)

        #Save total scaling constant and new mean of the whole pdf

        self.param_val[0] = 1.0/(self.x_max-self.x_min)*_mean

        self.param_val[1] = _mean * self.param_val[0]

        print(" pdf normailzed")

        print

    def param_list(self): #prints a list of all parameters of the form: number, name, value

        print("List of parameters")
        print(" Nr. Name Value")

        for i in range(len(self.param_val)):
            print(" {0}. {1}: {2}".format(i, self.param_str[i], self.param_val[i]))

        print

    def log_likelihood(self, x_part, cusp_amp): #Replaced soon with TMinuit

        _sum = 0.0

        for i in x_part:

            _sum += np.log(self.total_pdf(i**2, cusp_amp = cusp_amp, scan = True))

        return _sum
