#!/usr/bin/env python
# coding: utf-8

# # Import

# In[1]:


import numpy as np
from pdg_const import pdg
import matplotlib
import matplotlib.pyplot as plt
import pickle as pkl
import sys
import time
from helperfunctions import display_time, prepare_plot
import cmath as c
import scipy.integrate as integrate
from scipy.optimize import fminbound
from array import array as arr
import collections
from itertools import compress
import tensorflow as tf
import zfit
from zfit import ztf


# # Build model and graphs
# ## Create graphs

# In[2]:


def formfactor( 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

    mK = ztf.constant(pdg['Ks_M'])
    mbstar0 = ztf.constant(pdg["mbstar0"])
    mbstar = ztf.constant(pdg["mbstar"])
    b0 = ztf.constant(pdg["b0"])
    bplus = ztf.constant(pdg["bplus"])
    bT = ztf.constant(pdg["bT"])

    mmu = ztf.constant(pdg['muon_M'])
    mb = ztf.constant(pdg['bquark_M'])
    ms = ztf.constant(pdg['squark_M'])
    mB = ztf.constant(pdg['Bplus_M'])

    #N comes from derivation in paper

    N = 3

    #some helperfunctions

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

    z_oben = ztf.sqrt(tpos - q2) - ztf.sqrt(tpos - tzero)
    z_unten = ztf.sqrt(tpos - q2) + ztf.sqrt(tpos - tzero)
    z = tf.divide(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]*(tf.pow(z,i))

        return tf.complex(prefactor * _sum, ztf.constant(0.0))

    #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] * (tf.pow(z, i) - ((-1)**(i-N)) * (i/N) * tf.pow(z, N))

        return tf.complex(prefactor * _sum, ztf.constant(0.0))

def resonance(q, _mass, width, phase, scale):

    q2 = tf.pow(q, 2)

    mmu = ztf.constant(pdg['muon_M'])

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

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

    gamma_j = tf.divide(p, q2) * _mass * width / p0

    #Calculate the resonance

    _top = tf.complex(_mass * width, ztf.constant(0.0))

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

    com = _top/_bottom

    #Rotate by the phase

    r = tf.abs(com)

    _phase = tf.angle(com)

    _phase += phase

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

    com = tf.complex(scale* x, scale * y)

    return com

def bifur_gauss(q, mean, sigma_L, sigma_R, scale):

    _exp = tf.where(q < mean, ztf.exp(- tf.pow((q-mean),2) / (2 * sigma_L**2)), ztf.exp(- tf.pow((q-mean),2) / (2 * sigma_R**2)))

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

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

    com = ztf.complex(dgamma, ztf.constant(0.0))

    return com

def axiv_nonres(q):

    GF = ztf.constant(pdg['GF'])
    alpha_ew = ztf.constant(pdg['alpha_ew'])
    Vtb = ztf.constant(pdg['Vtb'])
    Vts = ztf.constant(pdg['Vts'])
    C10eff = ztf.constant(pdg['C10eff'])

    mmu = ztf.constant(pdg['muon_M'])
    mb = ztf.constant(pdg['bquark_M'])
    ms = ztf.constant(pdg['squark_M'])
    mK = ztf.constant(pdg['Ks_M'])
    mB = ztf.constant(pdg['Bplus_M'])

    q2 = tf.pow(q, 2)

    #Some helperfunctions

    beta = ztf.sqrt(tf.abs(1. - 4. * mmu**2. / q2))

    kabs = ztf.sqrt(mB**2. +tf.pow(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. * (tf.abs(Vtb*Vts))**2. * kabs * beta / (128. * np.pi**5.)

    #left term in bracket

    bracket_left = 2./3. * kabs**2. * beta**2. *tf.abs(tf.complex(C10eff, ztf.constant(0.0))*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 *tf.pow(tf.abs(tf.complex(C10eff, ztf.constant(0.0)) * formfactor(q2, "0")), 2)

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

    return prefactor1 * (bracket_left + bracket_middle) * 2 *ztf.sqrt(q2)

def vec(q, funcs):
    
    q2 = tf.pow(q, 2)

    GF = ztf.constant(pdg['GF'])
    alpha_ew = ztf.constant(pdg['alpha_ew'])
    Vtb = ztf.constant(pdg['Vtb'])
    Vts = ztf.constant(pdg['Vts'])
    C7eff = ztf.constant(pdg['C7eff'])

    mmu = ztf.constant(pdg['muon_M'])
    mb = ztf.constant(pdg['bquark_M'])
    ms = ztf.constant(pdg['squark_M'])
    mK = ztf.constant(pdg['Ks_M'])
    mB = ztf.constant(pdg['Bplus_M'])

    #Some helperfunctions

    beta = ztf.sqrt(tf.abs(1. - 4. * mmu**2. / q2))

    kabs = ztf.sqrt(mB**2. + tf.pow(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. * (tf.abs(Vtb*Vts))**2 * kabs * beta / (128. * np.pi**5.)

    #right term in bracket

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

    abs_bracket = tf.abs(c9eff(q, funcs) * formfactor(q2, "+") + tf.complex(2.0 * C7eff * (mb + ms)/(mB + mK), ztf.constant(0.0)) * formfactor(q2, "T"))**2

    bracket_right = prefactor2 * abs_bracket

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

    return prefactor1 * bracket_right * 2 * ztf.sqrt(q2)

def c9eff(q, funcs):

    C9eff_nr = tf.complex(ztf.constant(pdg['C9eff']), ztf.constant(0.0))

    c9 = C9eff_nr

    c9 = c9 + funcs

    return c9


# In[3]:


def G(y):
    
    def inner_rect_bracket(q):
        return tf.log(ztf.to_complex((1+tf.sqrt(q))/(1-tf.sqrt(q)))-tf.complex(ztf.constant(0), -1*ztf.constant(np.pi)))    
    
    def inner_right(q):
        return ztf.to_complex(2 * tf.atan(1/tf.sqrt(-q)))
    
    big_bracket = tf.where(y > ztf.const(0.0), inner_rect_bracket(y), inner_right(y))
    
    return ztf.to_complex(tf.sqrt(tf.abs(y))) * big_bracket

def h_S(m, q):
    
    return tf.constant(2) - G(tf.constant(1) - 4*tf.pow(m, 2) / tf.pow(q, 2))

def h_P(m,q):
    
    return 2/3 + (1 - (tf.constant(1) - 4*tf.pow(m, 2) / tf.pow(q, 2))) * h_S(m,q)


# ## Build pdf

# In[4]:


class total_pdf(zfit.pdf.ZPDF):
    _N_OBS = 1  # dimension, can be omitted
    _PARAMS = ['jpsi_mass', 'jpsi_scale', 'jpsi_phase', 'jpsi_width',
                'psi2s_mass', 'psi2s_scale', 'psi2s_phase', 'psi2s_width',
                'cusp_mass', 'sigma_L', 'sigma_R', 'cusp_scale'
                ]  # the name of the parameters

    def _unnormalized_pdf(self, x):
        
        x = x.unstack_x()

        def jpsi_res(q):
            return resonance(q, _mass = self.params['jpsi_mass'], scale = self.params['jpsi_scale'], phase = self.params['jpsi_phase'], width = self.params['jpsi_width'])

        def psi2s_res(q):
            return resonance(q, _mass = self.params['psi2s_mass'], scale = self.params['psi2s_scale'], phase = self.params['psi2s_phase'], width = self.params['psi2s_width'])

        def cusp(q):
            return bifur_gauss(q, mean = self.params['cusp_mass'], sigma_L = self.params['sigma_L'], sigma_R = self.params['sigma_R'], scale = self.params['cusp_scale'])

        funcs = jpsi_res(x) + psi2s_res(x) + cusp(x)

        vec_f = vec(x, funcs)

        axiv_nr = axiv_nonres(x)

        tot = vec_f + axiv_nr

        return tot


# ## Load data

# In[5]:


x_min = 2*pdg['muon_M']
x_max = (pdg["Bplus_M"]-pdg["Ks_M"]-0.1)

obs = zfit.Space('q', limits = (x_min, x_max))

with open(r"./data/slim_points/slim_points_toy_0_range({0}-{1}).pkl".format(int(x_min), int(x_max)), "rb") as input_file:
    part_set = pkl.load(input_file)

x_part = part_set['x_part']

x_part = x_part.astype('float64')

data = zfit.data.Data.from_numpy(array=x_part, obs=obs)


# ## Setup parameters

# In[6]:


#jpsi

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

jpsi_m = zfit.Parameter("jpsi_m", ztf.constant(jpsi_mass), floating = False)
jpsi_w = zfit.Parameter("jpsi_w", ztf.constant(jpsi_width), floating = False)
jpsi_p = zfit.Parameter("jpsi_p", ztf.constant(jpsi_phase), floating = False)
jpsi_s = zfit.Parameter("jpsi_s", ztf.constant(jpsi_scale), floating = False)

#psi2s

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

psi2s_m = zfit.Parameter("psi2s_m", ztf.constant(psi2s_mass), floating = False)
psi2s_w = zfit.Parameter("psi2s_w", ztf.constant(psi2s_width), floating = False)
psi2s_p = zfit.Parameter("psi2s_p", ztf.constant(psi2s_phase), floating = False)
psi2s_s = zfit.Parameter("psi2s_s", ztf.constant(psi2s_scale), floating = False)

#cusp

cusp_mass, sigma_R, sigma_L, cusp_scale = 3550, 3e-7, 200, 0

cusp_m = zfit.Parameter("cusp_m", ztf.constant(cusp_mass))
sig_L = zfit.Parameter("sig_L", ztf.constant(sigma_L))
sig_R = zfit.Parameter("sig_R", ztf.constant(sigma_R))
cusp_s = zfit.Parameter("cusp_s", ztf.constant(cusp_scale))


# ## Setup pdf

# In[7]:


total_f = total_pdf(obs=obs, jpsi_mass = jpsi_m, jpsi_scale = jpsi_s, jpsi_phase = jpsi_p, jpsi_width = jpsi_w,
            psi2s_mass = psi2s_m, psi2s_scale = psi2s_s, psi2s_phase = psi2s_p, psi2s_width = psi2s_w,
            cusp_mass = cusp_m, sigma_L = sig_L, sigma_R = sig_R, cusp_scale = cusp_s)

# print(total_pdf.obs)


# ## Test if graphs actually work and compute values

# In[8]:


def total_test_tf(xq):

    def jpsi_res(q):
        return resonance(q, jpsi_m, jpsi_s, jpsi_p, jpsi_w)

    def psi2s_res(q):
        return resonance(q, psi2s_m, psi2s_s, psi2s_p, psi2s_w)

    def cusp(q):
        return bifur_gauss(q, cusp_m, sig_L, sig_R, cusp_s)

    funcs = jpsi_res(xq) + psi2s_res(xq) + cusp(xq)

    vec_f = vec(xq, funcs)

    axiv_nr = axiv_nonres(xq)

    tot = vec_f + axiv_nr
    
    return tot

def jpsi_res(q):
    return resonance(q, jpsi_m, jpsi_s, jpsi_p, jpsi_w)

# calcs = zfit.run(total_test_tf(x_part))

test_q = np.linspace(x_min, x_max, 2000000)

probs = total_f.pdf(test_q)

calcs_test = zfit.run(probs)
res_y = zfit.run(jpsi_res(test_q))


# In[9]:


plt.clf()
# plt.plot(x_part, calcs, '.')
plt.plot(test_q, calcs_test, label = 'pdf')
# plt.plot(test_q, res_y, label = 'res')
plt.legend()
plt.ylim(0.0, 4e-4)
# plt.yscale('log')
# plt.xlim(3080, 3110)
plt.savefig('test.png')
print(jpsi_width)


# ## Adjust scaling of different parts

# In[10]:


# total_f.update_integration_options(draws_per_dim=10000000, mc_sampler=None)
# inte = total_f.integrate(limits = (3000, 3200), norm_range=False)
# print(zfit.run(inte))
# print(pdg["jpsi_BR"]/pdg["NR_BR"], zfit.run(inte)/pdg["NR_auc"])


# In[11]:


# factor_jpsi = pdg["NR_auc"]*pdg["jpsi_BR"]/(pdg["NR_BR"]*pdg["jpsi_auc"])
# print(np.sqrt(factor_jpsi))
# factor_psi2s = pdg["NR_auc"]*pdg["psi2s_BR"]/(pdg["NR_BR"]*pdg["psi2s_auc"])
# print(np.sqrt(factor_psi2s))


# In[12]:


# def _t_f(xq):

#     def jpsi_res(q):
#         return resonance(q, jpsi_m, jpsi_s, jpsi_p, jpsi_w)

#     def psi2s_res(q):
#         return resonance(q, psi2s_m, psi2s_s, psi2s_p, psi2s_w)

#     def cusp(q):
#         return bifur_gauss(q, cusp_m, sig_L, sig_R, cusp_s)

#     funcs = psi2s_res(xq) + jpsi_res(xq) + cusp(xq)

#     vec_f = vec(xq, funcs)

#     axiv_nr = axiv_nonres(xq)

#     tot = vec_f + axiv_nr
    
#     return tot

# def t_f(x):
#     probs = zfit.run(_t_f(ztf.constant(x)))
#     return probs


# In[13]:


print(36000*(1+ pdg["jpsi_BR"]/pdg["NR_BR"] + pdg["psi2s_BR"]/pdg["NR_BR"]))


# In[14]:


# start = time.time()

# result, err = integrate.quad(lambda x: t_f(x), x_min, x_max, limit = 50)
# print(result, "{0:.2f} %".format(err/result))
# print("Time:", time.time()-start)


# # Sampling
# ## One sample

# In[ ]:


nevents = 100

samp = total_f.sample(n=nevents)

sam = samp.unstack_x()

sam = zfit.run(sam)
# print(sam)


# In[ ]:


with open("data/zfit_toys/test_toy.pkl", "wb") as f:
    pkl.dump(sam, f, pkl.HIGHEST_PROTOCOL)


# In[ ]:


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

calcs = zfit.run(total_test_tf(samp))

plt.hist(sam, bins = bins, range = (x_min,x_max))

# plt.plot(sam, calcs, '.')
# plt.plot(test_q, calcs_test)
# plt.ylim(0, 0.0000007)
# plt.xlim(3000, 3750)

plt.savefig('test.png')


# ## Toys

# In[ ]:


# sampler = total_f.create_sampler(n=nevents)
# nll = zfit.loss.UnbinnedNLL(model=total_f, data=sampler, fit_range = (x_min, x_max))

# # for param in pdf.get_dependents():
# #     param.set_value(initial_value)

# sampler.resample(n=nevents)

# # Randomise initial values
# # for param in pdf.get_dependents():
# #     param.set_value(random value here)

# # Minimise the NLL
# minimizer = zfit.minimize.MinuitMinimizer(verbosity = 10)
# minimum = minimizer.minimize(nll)


# # Fitting

# In[ ]:


# nll = zfit.loss.UnbinnedNLL(model=total_f, data=data, fit_range = (x_min, x_max))

# minimizer = zfit.minimize.MinuitMinimizer()
# minimizer._use_tfgrad = False
# result = minimizer.minimize(nll)

# param_errors = result.error()

# for var, errors in param_errors.items():
#     print('{}: ^{{+{}}}_{{{}}}'.format(var.name, errors['upper'], errors['lower']))

# print("Function minimum:", result.fmin)


# In[ ]:


# samp = total_f.sample(n=nevents)


# In[ ]:


# sam = samp.unstack_x()

# sam = zfit.run(sam)

# bins = int((x_max-x_min)/7)

# calcs = zfit.run(total_test_tf(samp))

# plt.clf()

# plt.hist(sam, bins = bins, range = (x_min,x_max))

# # plt.plot(sam, calcs, '.')
# # plt.plot(test_q, calcs_test)
# # plt.ylim(0, 0.0000007)
# # plt.xlim(3000, 3750)

# plt.ylim(0,1000)

# plt.savefig('test.png')


# In[ ]: