from funclib.variables import *

#Funkce se bohuzel chova velmi nestabilne okolo polu M_minus. Hmotnost 
#M_eta byla zmenena na posledni chvili a graf se zacal chovam mnohem problematicteji  
#a musel jsem interpolovat pres vetsi cast. K2Lambda tedy bude obsahovat 
#vetsi chybu, ale vysledek z fitu se nezmeni nehlede na to, jestli pouzijeme
#orezanou nebo nestabilni verzi

M_eta = M_eta + epsj
M_minus = (M_eta- M_pi)**2 
M_plus = (M_eta + M_pi)**2

from funclib.k2tilde import K_2TildeReal
from funclib.k2tilde import K_2TildeImag

def K2Tilde(x):
    return K_2TildeReal(x) + 1j* K_2TildeImag(x)

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

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 - 4e-2, num=50)
energie3 = np.linspace(M_minus + 5e-2, 12*M_pi**2, num=50)

max_pocet = len(energie1)
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_2LambdaReal = [0]*3*max_pocet
k_2LambdaImag = [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_2LambdaReal[i] = real1[i] 
    k_2LambdaReal[i+max_pocet] = real2[i]
    k_2LambdaReal[i+2*max_pocet] = real3[i]
    k_2LambdaImag[i] = imag1[i] 
    k_2LambdaImag[i+max_pocet] = imag2[i]
    k_2LambdaImag[i+2*max_pocet] = imag3[i]
 

K_2LambdaReal = scipy.interpolate.interp1d(energie, k_2LambdaReal, fill_value='extrapolate')
K_2LambdaImag = scipy.interpolate.interp1d(energie, k_2LambdaImag, fill_value='extrapolate')

if __name__ == "__main__":

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