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 integrand(x):
    return 1/(16*(np.pi)) * M_pi**2/(sqrt(x*(x-4*M_pi**2))) *(1 + sqrt(x/(x-4*M_pi**2)) *  L(x) + M_pi**2/(x-4*M_pi**2) * (L(x))**2)

def numerical(x,s):
    return s/(16*np.pi**2)* 1/x * M_pi**2/(sqrt(x*(x-4*M_pi**2))) *(1 + sqrt(x/(x-4*M_pi**2)) *  L(x) + M_pi**2/(x-4*M_pi**2) * (L(x))**2) * 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))


from funclib.k1 import K_1Real
from funclib.k1 import K_1Imag

from funclib.k3 import K_3Real
from funclib.k3 import K_3Imag

def K1(x):
    return K_1Real(x) + 1j* K_1Imag(x)

def K3(x):
    return K_3Real(x) + 1j* K_3Imag(x)


def analytic1(x):
    return 1/3 * M_pi**2/(x-4*M_pi**2) * (K3(x) - x/(4*M_pi**2) * 1/32) + 1/2 * M_pi**2/(x-4*M_pi**2) * (K1(x) - x/(4*M_pi**2) * (-1/16)) + M_pi**2/(x-4*M_pi**2) * (Jbar1(x) - x/(4*M_pi**2) * Jbar2(4*M_pi**2))

def analytic2(x):
    return 1/3 * M_pi**2/(x-4*M_pi**2) * (K3(x) - x/(4*M_pi**2) * 1/32) + 1/2 * M_pi**2/(x-4*M_pi**2) * (K1(x) - x/(4*M_pi**2) * (-1/16)) + M_pi**2/(x-4*M_pi**2) * (Jbar2(x) - x/(4*M_pi**2) * Jbar2(4*M_pi**2))

energie_KLOE = np.loadtxt("funclib/hodnotytus.txt")
energie_ekvidistant = np.linspace(4*M_pi**2 + 6e-4, 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_41 = [0]*size
k_42 = [0]*size

real1 = [0]*size
real2 = [0]*size
imag1 = [0]*size
imag2 = [0]*size
Real1 = [0]*size
Real2 = [0]*size
Imag1 = [0]*size
Imag2 = [0]*size
k_4Real = [0]*2*size
k_4Imag = [0]*2*size

for i in range(size):
    NumReal1[i] = np.real(quad(numerical, 4*M_pi**2, 100000*M_pi**2, args=energie1[i], complex_func=True)[0]) 
    NumReal2[i] = np.real(quad(numerical, 4*M_pi**2, energie2[i]-eps, args=energie2[i], complex_func=True)[0] + quad(numerical, energie2[i]+eps, 100000*M_pi**2, args=energie2[i], complex_func=True)[0])
    NumImag1[i] = np.imag(quad(numerical, 4*M_pi**2, 100000*M_pi**2, args=energie1[i], complex_func=True)[0])
    NumImag2[i] = np.imag(quad(numerical, 4*M_pi**2, energie2[i]-eps, args=energie2[i], complex_func=True)[0] + quad(numerical, energie2[i]+eps, 100000*M_pi**2, args=energie2[i], complex_func=True)[0])
    Fi[i] = 1j * integrand(energie2[i])
    k_41[i] = NumReal1[i] + NumImag1[i] 
    k_42[i] = NumReal2[i] + NumImag2[i] + Fi[i]


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_4Real[i] = real1[i] 
    k_4Real[i+size] = real2[i]
    k_4Imag[i] = imag1[i] 
    k_4Imag[i+size] = imag2[i]


K_4Real = scipy.interpolate.interp1d(energie, k_4Real, fill_value='extrapolate')
K_4Imag = scipy.interpolate.interp1d(energie, k_4Imag, fill_value='extrapolate')


if __name__ == "__main__":
    

    plt.xlabel("s")
    plt.ylabel(r'$K_{4}(s)$')
    plt.plot(energie1, real1, color='orange', linestyle='--')
    plt.plot(energie2, real2, color='orange', linestyle='--')
    plt.plot(energie1, imag1, color='green')
    plt.plot(energie2, imag2, color='green')
    plt.plot(energie1, np.real(k_41), color='blue', linestyle='--')
    plt.plot(energie2, np.real(k_42), color='blue', linestyle='--')
    plt.plot(energie1, np.imag(k_41), color='red')
    plt.plot(energie2, np.imag(k_42), color='red')
    plt.show()
    
    
    srovnani1Real = [0]*size
    srovnani2Real = [0]*size
    srovnani1Imag = [0]*size
    srovnani2Imag = [0]*size
    for i in range(size):
        srovnani1Real[i] = real1[i] - np.real(k_41[i])
        srovnani2Real[i] = real2[i] - np.real(k_42[i])
        srovnani1Imag[i] = imag1[i] - np.imag(k_41[i])
        srovnani2Imag[i] = imag2[i] - np.imag(k_42[i])
        

    plt.xlabel("s")
    plt.plot(energie1, srovnani1Real, color='blue', linestyle='--')
    plt.plot(energie2, srovnani2Real, color='blue', linestyle='--')
    plt.plot(energie1, srovnani1Imag, color='red')
    plt.plot(energie2, srovnani2Imag, color='red')
    plt.show()