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

#Create pdfs

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, scale_axiv):

    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 scale_axiv * prefactor1 * (bracket_left + bracket_middle) * 2 *ztf.sqrt(q2)

def vec(q, scale_vec, funcs):

    scale = scale_vec

    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 scale_vec * 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

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',
                'scale_vec', 'scale_axiv',
                'cusp_mass', 'sigma_L', 'sigma_R', 'cusp_scale'
                ]  # the name of the parameters

    def _unnormalized_pdf(self, 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, self.params['scale_vec'], funcs)

        axiv_nr = axiv_nonres(x, self.params['scale_axiv'])

        tot = vec_f + axiv_nr

        return tot

#Load data

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)

#Build pdf

#vec/axiv scale

scale_v = zfit.Parameter("scale_vec", ztf.constant(1))
scale_a = zfit.Parameter("scale_axiv", ztf.constant(1))

#jpsi

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

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

#psi2s

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

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

#cusp

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

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))

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,
             scale_vec = scale_v, scale_axiv = scale_a,
            cusp_mass = cusp_m, sigma_L = sig_L, sigma_R = sig_R, cusp_scale = cusp_s)

# total_pdf = total_func.as_pdf()
print(total_pdf.obs)

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)
#_____________________________________________________________________

#
# class c9pdf(BaseFunctor):
#     __init__(c9pdf, param1, param2, param3, obs, name="c9pdf"):
#         params =
#
#   super()__init__(pdfs=c9pdf,...)
#
# def _unnormalized_pdf(..):
#     c9pdf = self.pdfs[0]
#     blabla... = bla  c9pdf.pdf(x)  ...
#
# #___________
#
# c9 = SumPDF(....)
# pdf = PDF(c9pdf=c9,...)
#
# #___________
#
# def _unnormalized_pdf(...):
#     a, b, c = x.unstack_x()
#     pdf1d = self.pdfs[0]
#     pdf1d.pdf(x)
#
# #___________
#
# c9pdf = self.pdfs[0]
# c9eff = c9pdf.pdf(x)