from funclib.variables import *

M_minus = (M_eta + epsj - M_pi)**2 
M_plus = (M_eta + epsj + M_pi)**2
deltaP = M_eta**2 + epsj - M_pi**2

def L(x):
    return np.emath.log((1-(sqrt((x-4*M_pi**2)/x)))/(1+(sqrt((x-4*M_pi**2)/x))))

def sigma1(x):
    return (1j* sqrt(-1 + 4*M_pi**2 / x))

def sigma2(x):
    return (sqrt(1 - 4*M_pi**2 / x))

def M(x):
    return -2*np.emath.log((1- deltaP/x + 1j*np.sqrt(-1*(x**2 + (M_eta**2 + epsj)**2 + M_pi**4 - 2*x*(M_eta**2 + epsj) - 2*x*M_pi**2 - 2*(M_eta**2 + epsj)*M_pi**2))/x)) + np.emath.log((4*M_pi**2)/x)

def kallen(x):
    return np.sqrt(M_plus - x +1j*eps1) * np.sqrt(M_minus - x +1j*eps1)

def kallen1(x):
    return -np.sqrt(M_plus - x) * np.sqrt(M_minus - x )

def kallen2(x):
    return np.sqrt(M_plus - x) * (1j*np.sqrt(-M_minus + x ))

def kallen3(x):
    return 1j*np.sqrt(-M_plus + x) * (1j*np.sqrt(-M_minus + x ))

def numerical_real1(x,s,z):
    return np.real(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

def numerical_real2(x,s,z):
    return np.real(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

def numerical_real3(x,s,z):
    return np.real(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

def numerical_imag1(x,s,z):
    return np.imag(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

def numerical_imag2(x,s,z):
    return np.imag(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

def numerical_imag3(x,s,z):
    return np.imag(-s**2/(16*(np.pi)**2) * 1/x**2 * L(x) * M(x) / kallen2(x) * 1/(x-s -z*1j)) *M_pi**2

eps  = 2.5e-6
eps1 = 2.5e-6

energie1 = np.linspace(0.1*M_pi**2, 4*M_pi**2 - 2e-5, num=50)
energie2 = np.linspace(4*M_pi**2 + 2e-5, M_minus - 2e-5, num=50)
energie3 = np.linspace(M_minus + 2e-5, 12*M_pi**2, num=50)


energie_KLOE = np.loadtxt("funclib/hodnotytus.txt")
body_energie2 = []
body_energie3 = []

for hodnota in energie_KLOE:
    if M_minus + 2e-10 <= hodnota <= 12 * M_pi**2:
        body_energie3.append(hodnota)
    elif 4 * M_pi**2 + 2e-10 <= hodnota <= M_minus - 2e-10:
        body_energie2.append(hodnota)

pocet_pridanych_energie3 = len(body_energie3)
pocet_pridanych_energie2 = len(body_energie2)

max_pocet = max(len(energie1), len(energie3) + pocet_pridanych_energie3, len(energie2) + pocet_pridanych_energie2)

energie1 = np.linspace(0.1*M_pi**2, 4*M_pi**2 - 2e-5, num=max_pocet)
energie2 = np.linspace(4*M_pi**2 + 2e-5, M_minus - 2e-5, num=(max_pocet - pocet_pridanych_energie2))
energie3 = np.linspace(M_minus + 2e-5, 12*M_pi**2, num=(max_pocet - pocet_pridanych_energie3))

energie1 = np.linspace(0.1*M_pi**2, 4*M_pi**2 - 2e-5, num=max_pocet)
energie2 = sorted(np.concatenate((energie2, body_energie2)))
energie3 = sorted(np.concatenate((energie3, body_energie3)))

energie1 = energie1[13:]
energie2 = energie2[:-13]
energie3 = energie3[13:]

max_pocet = max_pocet - 13

energie = [0]*3*max_pocet
for i in range(max_pocet):
    energie[i] = energie1[i]
    energie[i+max_pocet] = energie2[i]
    energie[i+2*max_pocet] = energie3[i]

real1 = [0]*max_pocet
real2 = [0]*max_pocet
real3 = [0]*max_pocet
imag1 = [0]*max_pocet
imag2 = [0]*max_pocet
imag3 = [0]*max_pocet
k_3DoubleTildeReal = [0]*3*max_pocet
k_3DoubleTildeImag = [0]*3*max_pocet

for i in range(max_pocet):
    real1[i] = quad(numerical_real1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie1[i], eps1))[0] + quad(numerical_real2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie1[i], eps1))[0] + quad(numerical_real3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie1[i], eps1))[0]
    real2[i] = quad(numerical_real1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie2[i], eps1))[0] + quad(numerical_real2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie2[i], eps1))[0] + quad(numerical_real3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie2[i], eps1))[0]
    real3[i] = quad(numerical_real1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie3[i], eps1))[0] + quad(numerical_real2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie3[i], eps1))[0] + quad(numerical_real3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie3[i], eps1))[0]
 


for i in range(max_pocet):
    imag1[i] = quad(numerical_imag1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie1[i], eps1))[0] + quad(numerical_imag2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie1[i], eps1))[0] + quad(numerical_imag3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie1[i], eps1))[0]
    imag2[i] = quad(numerical_imag1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie2[i], eps1))[0] + quad(numerical_imag2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie2[i], eps1))[0] + quad(numerical_imag3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie2[i], eps1))[0]
    imag3[i] = quad(numerical_imag1, 4*M_pi**2, np.real(M_minus) - eps, args=(energie3[i], eps1))[0] + quad(numerical_imag2, np.real(M_minus) + eps, np.real(M_plus) - eps, args=(energie3[i], eps1))[0] + quad(numerical_imag3, np.real(M_plus) + eps, 100000*M_pi**2, args=(energie3[i], eps1))[0]
  


for i in range(max_pocet):
    k_3DoubleTildeReal[i] = real1[i] 
    k_3DoubleTildeReal[i+max_pocet] = real2[i]
    k_3DoubleTildeReal[i+2*max_pocet] = real3[i]
    k_3DoubleTildeImag[i] = imag1[i] 
    k_3DoubleTildeImag[i+max_pocet] = imag2[i]
    k_3DoubleTildeImag[i+2*max_pocet] = imag3[i]

K_3DoubleTildeReal = scipy.interpolate.interp1d(energie, k_3DoubleTildeReal, fill_value='extrapolate')
K_3DoubleTildeImag = scipy.interpolate.interp1d(energie, k_3DoubleTildeImag, fill_value='extrapolate')


if __name__ == "__main__":


    plt.xlabel("s")
    plt.ylabel(r'$\tilde{\tilde{K}}_{3}(s)$')
    plt.plot(energie, k_3DoubleTildeReal, color='orange', linestyle='--')
    plt.plot(energie, k_3DoubleTildeImag, color='green')
    plt.show() 
