from funclib.variables import *
from getting_op import E0_pm_Op4

from funclib.jBarF import JBarReal
from funclib.jBarF import JBarImag
#F_pi = sympy.symbols("F_pi")
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)    

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) + E0_pm_Op4)**2
temp = full_expression.subs(A_x, A_x/(F_pi**2))
expression = temp.expand()

filtered_terms = []

terms = sympy.Add.make_args(expression)

for term in terms:
    if 'F_pi' in str(term): 
        if '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_terms.append(term)

final_expression0 = sympy.Add(*filtered_terms)      

op4 = final_expression0

def op4_eta(X, A_x, B_x, C_x, D_x):
    t, u, s = X
    return (abs(eval(str(op4))))

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

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

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

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

full_expression_2 = (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) + E0_pm_Op4)
temp2 = full_expression_2.subs(A_x, A_x/(F_pi**2))
expression = temp2.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_terms0 = filtered_terms0
final_expression0 = sympy.Add(*final_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)


op4 = final_expression0 * (final_expression0 + final_expression1 + final_expression2 + final_expression3) + final_expression1 * (final_expression0 + final_expression1 + final_expression2) + final_expression2 * (final_expression0 + final_expression1) + final_expression3 * (final_expression0)


def op4_eta(X, A_x, B_x, C_x, D_x):
    t, u, s = X
    return (abs(eval(str(op4))))

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

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

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

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


