from getting_op import E0_pm_Op6
from funclib.variables import *


from funclib.jBarF import JBarReal
from funclib.jBarF import JBarImag


from funclib.k0 import K_0Real 
from funclib.k0 import K_0Imag

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

from funclib.k0doubletilde import K_0DoubleTildeReal
from funclib.k0doubletilde import K_0DoubleTildeImag

from funclib.k0lambda import K_0LambdaReal
from funclib.k0lambda import K_0LambdaImag



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

from funclib.k1tilde import K_1TildeReal
from funclib.k1tilde import K_1TildeImag

from funclib.k1doubletilde import K_1DoubleTildeReal
from funclib.k1doubletilde import K_1DoubleTildeImag

from funclib.k1lambda import K_1LambdaReal
from funclib.k1lambda import K_1LambdaImag



from funclib.k2 import K_2Real
from funclib.k2 import K_2Imag

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

from funclib.k2doubletilde import K_2DoubleTildeReal
from funclib.k2doubletilde import K_2DoubleTildeImag

from funclib.k2lambda import K_2LambdaReal
from funclib.k2lambda import K_2LambdaImag



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

from funclib.k3tilde import K_3TildeReal
from funclib.k3tilde import K_3TildeImag

from funclib.k3doubletilde import K_3DoubleTildeReal
from funclib.k3doubletilde import K_3DoubleTildeImag

from funclib.k3lambda import K_3LambdaReal
from funclib.k3lambda import K_3LambdaImag

from funclib.k4 import K_4Real
from funclib.k4 import K_4Imag

hodnoty = np.loadtxt("funclib/hodnoty.txt")
t_data = (hodnoty[:,0])
u_data = (hodnoty[:,1])
s_data = (hodnoty[:,2])
func_data = (hodnoty[:,3])
sizeS = len(s_data)

hodnoty_sloupec = np.loadtxt("funclib/hodnotytus.txt")
size = len(hodnoty_sloupec)

def JBar(x, M_pi):
    return JBarReal(x) + 1j*JBarImag(x)


def KBar_0(x, M_pi):
    return K_0Real(x) + 1j*K_0Imag(x)
def KTilde_0(x, M_pi):
    return K_0TildeReal(x) + 1j*K_0TildeImag(x)
def KTildeTilde_0(x, M_pi):
    return K_0DoubleTildeReal(x) + 1j*K_0DoubleTildeImag(x)

def KLambda_0(x, M_pi, M_eta):
    return K_0LambdaReal(x) + 1j*K_0LambdaImag(x)
def KBar_1(x, M_pi):
    return K_1Real(x) + 1j*K_1Imag(x)
def KTilde_1(x, M_pi):
    return K_1TildeReal(x) + 1j*K_1TildeImag(x)
def KTildeTilde_1(x, M_pi):
    return K_1TildeReal(x) + 1j*K_1TildeImag(x)
def KLambda_1(x, M_pi, M_eta):
    return K_1LambdaReal(x) + 1j*K_1LambdaImag(x)

def KBar_2(x, M_pi):
    return K_2Real(x) + 1j*K_2Imag(x)
def KTilde_2(x, M_pi, M_eta):
    return K_2TildeReal(x) + 1j*K_2TildeImag(x)
def KTildeBar_2(x, M_pi, M_eta):
    return K_2TildeReal(x) + 1j*K_2TildeImag(x)
def KTildeTilde_2(x, M_pi, M_eta):
    return K_2DoubleTildeReal(x) + 1j*K_2DoubleTildeImag(x)
def KLambda_2(x, M_pi, M_eta):
    return K_2LambdaReal(x) + 1j*K_2LambdaImag(x)

def KBar_3(x, M_pi):
    return K_3Real(x) + 1j*K_3Imag(x)
def KTilde_3(x, M_pi, M_eta):
    return K_3TildeReal(x) + 1j*K_3TildeImag(x)
def KTildeTilde_3(x, M_pi, M_eta):
    return K_3TildeReal(x) + 1j*K_3TildeImag(x)
def KLambda_3(x, M_pi, M_eta):
    return K_3LambdaReal(x) + 1j*K_3LambdaImag(x)

def KBar_4(x, M_pi):
    return K_4Real(x) + 1j*K_4Imag(x)

full_expression = (A_x + B_x/(F_pi**2) *(s - s_0) + C_x/(F_pi**4) *(s - s_0)**2 + D_x/(F_pi**4) *((t - s_0)**2 + (u - s_0)**2) + (E_x/(F_pi**6) * (s - s_0)**3 + F_x/(F_pi**6) * ((t - s_0)**3 + (u - s_0)**3))+ E0_pm_Op6)
temp = full_expression.subs(A_x, A_x/(F_pi**2))
expression = temp.expand()
terms = sympy.Add.make_args(expression)

filtered_terms0 = []

for term in terms:
    if 'F_pi' in str(term): 
        if 'F_pi**1' in str(term) or 'F_pi**2' in str(term) or 'F_pi**3' in str(term) or 'F_pi**4' in str(term) or 'F_pi**5' in str(term) or 'F_pi**6' in str(term) or 'F_pi**7' in str(term) or 'F_pi**8' in str(term) or 'F_pi**9' in str(term) or 'F_pi**10' in str(term) or 'F_pi**11' in str(term) or 'F_pi**12' in str(term):
            continue  
    filtered_terms0.append(term)
final_expression0 = sympy.Add(*filtered_terms0)


filtered_terms1 = []

for term in terms:
    if 'F_pi' in str(term): 
        if 'F_pi**3' in str(term) or 'F_pi**4' in str(term) or 'F_pi**5' in str(term) or 'F_pi**6' in str(term) or 'F_pi**7' in str(term) or 'F_pi**8' in str(term) or 'F_pi**9' in str(term) or 'F_pi**10' in str(term) or 'F_pi**11' in str(term) or 'F_pi**12' in str(term):
            continue  
    filtered_terms1.append(term)

final_terms1 = []
final_terms1 = [term for term in filtered_terms1 if 'F_pi' in str(term)]
final_expression1 = sympy.Add(*final_terms1)


filtered_terms2 = []

for term in terms:
    if 'F_pi' in str(term): 
        if 'F_pi**1' in str(term) or 'F_pi**2' in str(term) or 'F_pi**5' in str(term) or 'F_pi**6' in str(term) or 'F_pi**7' in str(term) or 'F_pi**8' in str(term) or 'F_pi**9' in str(term) or 'F_pi**10' in str(term) or 'F_pi**11' in str(term) or 'F_pi**12' in str(term):
            continue  
    filtered_terms2.append(term)

final_terms2 = []
final_terms2 = [term for term in filtered_terms2 if 'F_pi' in str(term)]
final_expression2 = sympy.Add(*final_terms2)


filtered_terms3 = []

for term in terms:
    if 'F_pi' in str(term): 
        if 'F_pi**1' in str(term) or 'F_pi**2' in str(term) or 'F_pi**3' in str(term) or 'F_pi**4' in str(term) or 'F_pi**7' in str(term) or 'F_pi**8' in str(term) or 'F_pi**9' in str(term) or 'F_pi**10' in str(term) or 'F_pi**11' in str(term) or 'F_pi**12' in str(term):
            continue  
    filtered_terms3.append(term)

final_terms3 = []
final_terms3 = [term for term in filtered_terms3 if 'F_pi' in str(term)]
final_expression3 = sympy.Add(*final_terms3)


filtered_terms4 = []

for term in terms:
    if 'F_pi' in str(term): 
        if 'F_pi**1' in str(term) or 'F_pi**2' in str(term) or 'F_pi**3' in str(term) or 'F_pi**4' in str(term) or 'F_pi**5' in str(term) or 'F_pi**6' in str(term) or 'F_pi**9' in str(term) or 'F_pi**10' in str(term) or 'F_pi**11' in str(term) or 'F_pi**12' in str(term):
            continue  
    filtered_terms4.append(term)

final_terms4 = []
final_terms4 = [term for term in filtered_terms4 if 'F_pi' in str(term)]
final_expression4 = sympy.Add(*final_terms4)


op6 = final_expression0 * (final_expression0 + final_expression1 + final_expression2 + final_expression3 + final_expression4) + final_expression1 * (final_expression0 + final_expression1 + final_expression2 + final_expression3) + final_expression2 * (final_expression0 + final_expression1 + final_expression2) + final_expression3 * (final_expression0 + final_expression1) + final_expression4 * final_expression0

def op6_eta(X, A_x, B_x, C_x, D_x, E_x, F_x):
    t, u, s = X
    return (abs(eval(str(op6))))

popt, pcov = curve_fit(op6_eta,(t_data, u_data, s_data), func_data)


print("A_x, B_x, C_x, D_x, E_x, F_x: " + str(popt))

perr = np.sqrt(np.diag(pcov))
print("Error: " + str(perr))

cor = np.zeros((6,6))
for i in range (6):
    for j in range (6):
     cor[i,j] = abs(pcov[i,j] * 1/(np.sqrt(pcov[i,i] * pcov[j,j])))
          
print("Correlation matrix")    
print(cor)


























