import numpy as np
import matplotlib.pyplot as plt
from scipy.special import roots_legendre
import cmath
from time import time

# constant values (GeV)

M_roh = 0.7736
M_pi = 0.13957
s_0 = 4*M_pi**2
M_K = 0.496
B_0 = 1.055
B_1 = 0.15
lambda_1 = 1.57
lambda_2 = -1.96
lambda_high = 10
epsilon = 0.000000001
s_mid_high = 1.42**2 # this is the s value, so it is in GeV^2

# derivative values
lambda_0 = 0 # this is fixed by the low energy parametrization: lambda_0 = delta1_low(4M_K**2) (sqrt(s)=2M_K)
delta1_mid_high = 0 # again, this is fixed by the value of delta1_mid at 1.42 GeV (see later in the code)

# Numerical constants
gl_n_omnes = 100 # number of points used in the Gauss Legendre Integration
gl_n_phi0 = 100

gl_lookup = {}
omnes_lookup = {}
delta_lookup = {}

# Print iterations progress
def printProgressBar (iteration, total, prefix = '', suffix = '', decimals = 1, length = 100, fill = '█', printEnd = "\r"):
    """
    Call in a loop to create terminal progress bar
    @params:
        iteration   - Required  : current iteration (Int)
        total       - Required  : total iterations (Int)
        prefix      - Optional  : prefix string (Str)
        suffix      - Optional  : suffix string (Str)
        decimals    - Optional  : positive number of decimals in percent complete (Int)
        length      - Optional  : character length of bar (Int)
        fill        - Optional  : bar fill character (Str)
        printEnd    - Optional  : end character (e.g. "\r", "\r\n") (Str)
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filledLength = int(length * iteration // total)
    bar = fill * filledLength + '-' * (length - filledLength)
    print(f'\r{prefix} |{bar}| {percent}% {suffix}', end = printEnd)
    # Print New Line on Complete
    if iteration == total: 
        print()

''' the standard implementation of arctan gives angles in (-pi/2, pi/2). This gives angles in (0, pi). Note: This doesn't imply a simple shift of all values by pi, but a transposition of the results from (-pi/2, 0) to (pi/2, pi).'''
def arctan_shift(x):
	atan_std = np.arctan(x)
	if atan_std < 0:
		return atan_std + np.pi
	return atan_std
	
def log_shift(x):
	log_std = cmath.log(x)
	if log_std.imag < 0:
		return 2*np.pi*1j + log_std
	else:
		return log_std

def omega(s):
    # in this instance, Elena uses s_0**(1/2) = 1.05 GeV, but later she uses s_0 = 4*M_pi**2
	a = np.sqrt(s) 
	b = np.sqrt(1.05**2 - s)
	return (a - b) / (a + b)
	

# ----------------
# Delta function section
# ----------------

def k(s):
	return np.sqrt(s/4-M_pi**2)

def delta1_low(s):
    #a = np.sqrt(s) / (2*k(s)**3) * (M_roh**2 - s) * ( (2*M_pi**3) / (M_roh**2*np.sqrt(s)) + B_0 + B_1*omega(s) )
    return arctan_shift((np.sqrt(s) / (2*k(s)**3) * (M_roh**2 - s) * ( (2*M_pi**3)/(M_roh**2*np.sqrt(s)) + B_0 + B_1*omega(s)))**-1)

def delta1_mid(s):
	return lambda_0 + lambda_1 * (0.5*np.sqrt(s) / M_K - 1) + lambda_2 * (0.5*np.sqrt(s) / M_K - 1)**2

def delta1_high(s):
	return np.pi + (delta1_mid_high - np.pi) * (lambda_high**2 + s_mid_high) / (lambda_high**2 + s)
	

def delta(s):
	# if s in delta_lookup:
		# return delta_lookup[s]
	# if s < (2*M_K)**2:
		# delta_lookup[s] = delta1_low(s)
	# elif s < s_mid_high:
		# delta_lookup[s] = delta1_mid(s)
	# else:
		# delta_lookup[s] = delta1_high(s)
	# return delta_lookup[s]
	if s < (2*M_K)**2:
		return delta1_low(s)
	elif s < s_mid_high:
		return delta1_mid(s)
	else:
		return delta1_high(s)
		
		
# set cutoff values
lambda_0 = delta1_low((2*M_K)**2)
delta1_mid_high = delta1_mid(s_mid_high)
		
		
# ----------------
# Omnes Integral section
# ----------------

delta_time = 0

def omnes_integrand(s_tick,s):
    global delta_time
    delta_s_tick = delta(s_tick)
    delta_s = delta(s)
    return (delta_s_tick - delta_s)/(s_tick*(s_tick-s+epsilon*1j))


def integral_gl(f,s,a,b,n, prog=False):
	if n not in gl_lookup:
		# print("now calculating roots, weights")
		t0 = time()
		gl_lookup[n]  = roots_legendre(n)
		#print("time to calculate roots and weights with n={}: {} seconds".format(n,time() - t0))
	roots, weights = gl_lookup[n]
	sum = 0
	for i in range(n):
		if prog:
			printProgressBar(i,n)
		sum += weights[i]*f(roots[i]*(b-a)/2 + (a+b)/2,s)
	sum *= (b-a)/2
	return sum

def subst(f, phi, phi_deriv):
    return lambda x,s : f(phi(x),s)*phi_deriv(x)

def integral_gl_tan_reparam(f,s,a,b,n, prog=False):
    return integral_gl(subst(f,np.tan,lambda x : (1 / np.cos(x)**2)),s,np.arctan(a),np.pi / 2,n,prog)

def omnes(s, a=s_0, inf = 10000000000, optimized = True, n = gl_n_omnes):
	if s in omnes_lookup:
		return omnes_lookup[s]
	elif optimized:
		analyt = delta(s) / s * cmath.log(((inf-s-epsilon*1j)*a) / ((a-s-epsilon*1j)*inf))
		res = cmath.exp(s/np.pi*(integral_gl_tan_reparam(omnes_integrand,s,a,inf,n) + analyt))
		#omnes_lookup[s] = res
		return res
	else:
		res = cmath.exp(s/np.pi*(integral_gl(omnes_integrand,s,a,inf,n)+delta(s) / s * cmath.log(((inf-s-epsilon*1j)*a) / ((a-s-epsilon*1j)*inf))))
		#omnes_lookup[s] = res
		return res
		# A more easily readible version of the formula
		# if optimized == True:  
		# 	m = integral_gl_tan_reparam(omnes_integrand,s,a,inf,n)-
		# else: 
		# 	m = integral_gl(omnes_integrand,s,a,inf,n)
		# n = ((inf-s-epsilon*1j)*a) / ((a-s-epsilon*1j)*inf)
		# p = delta(s) / s * cmath.log(n)
		# omnes_lookup[s] = cmath.exp(s/np.pi*(m+p))
		# if s > 0.74**2 and s < 0.76**2:
		#	print(s, m, n, p, cmath.exp(s/np.pi*(m+p)))
		# return cmath.exp(s/np.pi*(m+p))

# ----------------
# Phase reconstruction section
# ----------------

def phi0_integrand_bad(s_tick, s): 
	omnes_s = omnes(s)
	omnes_s_tick = omnes(s_tick)
	#a = np.log(cmath.polar(omnes_s_tick / omnes_s)[0]**2)
	#b = np.sqrt(s_tick - s_0) * s_tick * (s_tick - s)
	return np.log(cmath.polar(omnes_s_tick / omnes_s)[0]**2) / (np.sqrt(s_tick - s_0) * s_tick * (s_tick - s))

def phi0_integrand(x, s):
	omnes_s = omnes(s)
	omnes_x = omnes(x**2 + s_0)
	
	#a = np.log(cmath.polar(omnes_s_tick / omnes_s)[0]**2)
	#b = np.sqrt(s_tick - s_0) * s_tick * (s_tick - s)
	return 2*np.log(cmath.polar(omnes_x / omnes_s)[0]**2) / ((s_0 + x**2) * (s_0 + x**2 - s))

def phi0(s, a=0, inf = 100000, optimized = True, n = gl_n_phi0):
	c = np.pi / (s * np.sqrt(s_0)) * np.log(cmath.polar(omnes(s))[0]**2)
	
	if optimized == True:
		integral = integral_gl_tan_reparam(phi0_integrand, s, a, inf, n, False)
	else:
		integral = integral_gl(phi0_integrand, s, a, inf, n)
	return -s * np.sqrt(s-s_0) / (2*np.pi) * (integral - c)

# ----------------
# Plotting section
# ----------------


fig, ax = plt.subplots()

# these are sqrt(s) values, not s
x = np.linspace(s_0**0.5+0.001,3,500)
#x = np.linspace(0.7, 0.85, 500)
x_to_pi2 = np.linspace(np.arctan(s_0)**0.5+0.001,(np.pi/2)**0.5,5000)

y_delta = [delta(i**2) for i in x]
#y1 = [cmath.polar(omnes_integrand(i**2,0.73))[0]**2 for i in x]
#y2 = [cmath.polar(omnes_integrand(i**2,0.73965))[0]**2 for i in x]
#y3 = [cmath.polar(omnes_integrand(i**2,0.75))[0]**2 for i in x]
#y_omnes_integrand_tan = [cmath.polar(omnes_integrand(np.tan(i**2),2) / np.cos(i**2)**2)[0]**2 for i in x_to_pi)]
#y_omnes_acc = [cmath.polar(omnes(i**2, optimized=True))[0]**2 for i in x]
#y_omnes_best = [cmath.polar(omnes(i**2, inf=10000, n = 10000))[0]**2 for i in x]
#y_phi0_integrand = [phi0_integrand(i**2, 2) for i in x]
#y_phi0_integrand_tan = [phi0_integrand(np.tan(i**2), 0.8)*1/np.cos(i**2)**2 for i in x_to_pi2]
#y_phi0_integrand_tan2 = [phi0_integrand(np.tan(i**2), 10000)*1/np.cos(i**2)**2 for i in x_to_pi2]
#y_phi0_integrand2 = [phi0_integrand(i**2, 0.8035457451490298) for i in x_phi0]

'''
for n in [10,100,1000,10000]:
	print("Phi0 for n_omnes={}, n_phi0={}: {}".format(gl_n_omnes, n, phi0(100**2, n=n)))
'''

'''
n_values = [10,100,1000,10000,100000]
y_phi0 = [[] for _ in n_values]
for (i_n,n) in enumerate(n_values):
	
# 	for (i,val) in enumerate(x):
# 		y_phi0[i_n].append(phi0(val**2, n = n))
# 		printProgressBar(i,len(x))
	print("Ready with {}".format(n))
	print()
'''
# Test sensitivity of phi0 to n compared to omnes
'''
n_values = np.arange(5, 4000, 25)
y_omnes_n_cmp = [cmath.polar(omnes(0.8**2, n=i))[0]**2 for i in n_values]
y_phi0_n_cmp = [phi0(0.8**2, n=i) for i in n_values]
ax.plot(n_values, y_omnes_n_cmp, label = "omnes")
ax.plot(n_values, y_phi0_n_cmp, label = "phi0")
'''

ax.plot(x, y_delta, label = r"$\delta$")
#ax.plot(x, y2, label = "delta")
#ax.plot(x, y3, label = "delta")
#ax.plot(x_to_pi2, y_omnes_integrand_tan, label = r"$\Omega$ integrand w/ tan-param.")
#ax.plot(x, y_omnes, label = r'${|\Omega (s)|}^2$')
#ax.plot(x, y_omnes_acc, label = r'${|\Omega (s)|}^2$, reparam')
#ax.plot(x, y_omnes_best, label = r'${|\Omega (s)|}^2$, but best')
#ax.plot(x, y_phi0_integrand, label = r'$\phi_0$ Integrand')
#ax.plot(x_to_pi2, y_phi0_integrand_tan, label = r'$\phi_0$ integrand w/ tan-param.')
#ax.plot(x_to_pi2, y_phi0_integrand_tan2, label = r'$\phi_0$ integrand w/ tan-param, s=10000.')
#ax.plot(x_phi0, y_phi0_integrand2, label = r'$\phi_0$ Integrand2')
#for (i,y) in enumerate(y_phi0):
	#ax.plot(x, y, label = r'$\phi_0, n={}$'.format(i))

# t0 = time()
# y_phi0 = [phi0(i**2) for i in x]
# ax.plot(x, y_phi0)
# print(time() - t0)

t0 = time()
print(phi0(10000, n=100))
print(time() - t0)

# ax.legend()
# #plt.yscale('log')
# ax.grid(which='major', color='#A4A4A4', linewidth=1)
# ax.grid(which='minor', color='#B5B5B5', linestyle=':', linewidth=0.5)
# ax.minorticks_on()
# plt.show()