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 integrand(x):
    return (1/(16*np.pi)* sigma2(x)) 

def numerical(x,s):
    return (s/(16*np.pi**2) * 1/x * sigma2(x)) * 1/(x-s)

def Jbar1(x):
    return 1/(16*np.pi**2) * (2 - 2*np.sqrt((4*M_pi**2 - x)/(x)) * np.arctan(np.sqrt((x)/(4*M_pi**2-x))))

def Jbar2(x):
    return 1/(16*np.pi**2) * (2 + sigma(x)* (np.emath.log((1-sigma(x))/(1+sigma(x))) + 1j*np.pi))

def analytic1(x):
    return Jbar1(x) 

def analytic2(x):
    return Jbar2(x)

energie_KLOE = np.loadtxt("funclib/hodnotytus.txt")
energie_ekvidistant = np.linspace(4*M_pi**2 + 5e-6, 12*M_pi**2, num = 50)
energie2 = np.sort( np.concatenate([energie_KLOE, energie_ekvidistant]))

size = len(energie2)

energie1 = np.linspace(0.1*M_pi**2, 4*M_pi**2 - 5e-6, num = size)

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

NumReal1 = [0]*size
NumReal2 = [0]*size
NumImag1 = [0]*size 
NumImag2 = [0]*size 
Fi = [0]*size 
k_01 = [0]*size
k_02 = [0]*size

real1 = [0]*size
real2 = [0]*size
imag1 = [0]*size
imag2 = [0]*size
k_0TildeReal = [0]*2*size
k_0TildeImag = [0]*2*size


for i in range(size):
    real1[i] = np.real(analytic1(energie1[i]))
    real2[i] = np.real(analytic2(energie2[i]))
    imag1[i] = np.imag(analytic1(energie1[i]))
    imag2[i] = np.imag(analytic2(energie2[i]))
    k_0TildeReal[i] = real1[i] 
    k_0TildeReal[i+size] = real2[i]
    k_0TildeImag[i] = imag1[i] 
    k_0TildeImag[i+size] = imag2[i]


K_0TildeReal = scipy.interpolate.interp1d(energie, k_0TildeReal, fill_value='extrapolate')
K_0TildeImag = scipy.interpolate.interp1d(energie, k_0TildeImag, fill_value='extrapolate')


if __name__ == "__main__":

    
    plt.xlabel("s")
    plt.ylabel(r'$\tilde{K}_{0}(s)$')
    plt.plot(energie1, real1, color='blue', linestyle='--')
    plt.plot(energie2, real2, color='blue', linestyle='--')
    plt.plot(energie1, imag1, color='red')
    plt.plot(energie2, imag2, color='red')
    plt.show()
    
    