from funclib.variables import *

def sigma(x):
    return (sqrt(1 - 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 numeric_real1(x,s,z):
    return np.real((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

def numeric_real2(x,s,z):
    return np.real((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

def numeric_real3(x,s,z):
    return np.real((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

def numeric_imag1(x,s,z):
    return np.imag((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

def numeric_imag2(x,s,z):
    return np.imag((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

def numeric_imag3(x,s,z):
    return np.imag((s**2/(16*np.pi**2) * 1/x**2 * sigma2(x)) * 1/(x-s -z*1j))

eps  = 1e-5
eps1 = 2e-5

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-5 <= hodnota <= 12 * M_pi**2:
        body_energie3.append(hodnota)
    elif 4 * M_pi**2 + 2e-5 <= hodnota <= M_minus - 2e-5:
        body_energie2.append(hodnota)

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

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[5:]
energie2 = energie2[:-5]
energie3 = energie3[5:]

max_pocet = max_pocet - 5

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_0DoubleTildeReal = [0]*3*max_pocet
k_0DoubleTildeImag = [0]*3*max_pocet

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


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

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

K_0DoubleTildeReal = scipy.interpolate.interp1d(energie, k_0DoubleTildeReal, fill_value='extrapolate')
K_0DoubleTildeImag = scipy.interpolate.interp1d(energie, k_0DoubleTildeImag, fill_value='extrapolate')


if __name__ == "__main__":


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

