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(abs(1- 4*M_pi**2 / x)))

def L(x):
    return np.emath.log((1-sigma(x))/(1+sigma(x)))

def L1(x):
    return np.emath.log((1-sigma1(x))/(1+sigma1(x)))

def L2(x):
    return np.emath.log(abs((1-sigma2(x))/(1+sigma2(x))))

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

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

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

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

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

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

eps  = 0.6e-5
eps1 = 3e-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-6, num=50)
energie3 = np.linspace(M_minus + 2e-6, 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_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[3:]
energie2 = energie2[:-3]
energie3 = energie3[3:]

max_pocet = max_pocet - 3

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_1DoubleTildeReal = [0]*3*max_pocet
k_1DoubleTildeImag = [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_1DoubleTildeReal[i] = real1[i] 
    k_1DoubleTildeReal[i+max_pocet] = real2[i]
    k_1DoubleTildeReal[i+2*max_pocet] = real3[i]
    k_1DoubleTildeImag[i] = imag1[i] 
    k_1DoubleTildeImag[i+max_pocet] = imag2[i]
    k_1DoubleTildeImag[i+2*max_pocet] = imag3[i]

K_1DoubleTildeReal = scipy.interpolate.interp1d(energie, k_1DoubleTildeReal, fill_value='extrapolate')
K_1DoubleTildeImag = scipy.interpolate.interp1d(energie, k_1DoubleTildeImag, fill_value='extrapolate')


if __name__ == "__main__":

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

