from funclib.variables import *


from funclib.k0tilde import K_0TildeReal
from funclib.k0tilde import K_0TildeImag

def K0Tilde(x):
    return K_0TildeReal(x) + 1j* K_0TildeImag(x)


def analytic(x):
    term1 = (M_pi**2 / (x - M_plus)) * (K0Tilde(x) - (x / M_plus) * K0Tilde(M_plus))
    term2 = (M_pi**2 / (x - M_minus)) * (K0Tilde(x) - (x / M_minus) * K0Tilde(M_minus))
    return (1 / 4) * (term1 - term2)

energie1 = np.linspace(0.1*M_pi**2, 4*M_pi**2 - 2e-6, num=50)
energie2 = np.linspace(4*M_pi**2 + 2e-6, 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-6, num=max_pocet)
energie2 = np.linspace(4*M_pi**2 + 2e-6, M_minus - 2e-6, num=(max_pocet - pocet_pridanych_energie2))
energie3 = np.linspace(M_minus + 2e-6, 12*M_pi**2, num=(max_pocet - pocet_pridanych_energie3))

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

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_0LambdaReal = [0]*3*max_pocet
k_0LambdaImag = [0]*3*max_pocet


for i in range(max_pocet):
    real1[i] = np.real(analytic(energie1[i]))
    real2[i] = np.real(analytic(energie2[i]))
    real3[i] = np.real(analytic(energie3[i]))
    imag1[i] = np.imag(analytic(energie1[i]))
    imag2[i] = np.imag(analytic(energie2[i]))
    imag3[i] = np.imag(analytic(energie3[i]))
    k_0LambdaReal[i] = real1[i] 
    k_0LambdaReal[i+max_pocet] = real2[i]
    k_0LambdaReal[i+2*max_pocet] = real3[i]
    k_0LambdaImag[i] = imag1[i] 
    k_0LambdaImag[i+max_pocet] = imag2[i]
    k_0LambdaImag[i+2*max_pocet] = imag3[i]
 


K_0LambdaReal = scipy.interpolate.interp1d(energie, k_0LambdaReal, fill_value='extrapolate')
K_0LambdaImag = scipy.interpolate.interp1d(energie, k_0LambdaImag, fill_value='extrapolate')

if __name__ == "__main__":


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

