reset()
var('z')
var('psi')

def reset_globals():
	global omega0, omega1, omega2, omega1h, omega2h, t, c
	global varpi1, varpi2, c1, c2
	omega0(z) = 1
	omega1h(z) = ln(z) - 6*z
	omega2h(z) = ln(z)^2 - 12*ln(z)*z-18*z
	c = {}
	c[1] = -18
	t = {}
	t[1] = -6
	varpi1(psi) = 0
	varpi2(psi) = 0
	c1 = {}
	c2 = {}

def theta(f):
	return z*diff(f, z)

def diff_op_z(f):
	g = theta(f)
	g = 3*theta(g) + g
	g = 3*theta(g) + 2*g
	return theta(theta(theta(f))) + 3*z*g
# ToDo: Need the same in turns of psi!!!
	

def theta_psi(f):
	return psi*diff(f, psi)

def diff_op_psi(f):
	g = -theta_psi(f) + f
	g = -theta_psi(g) + 2*g
	g = theta_psi(g)
	return -theta_psi(theta_psi(theta_psi(f))) + 1/psi^3*g
	
# This function initializes the power series expansions around z=0
# up to order N
def set_power_series0(N):
	global omega0, omega1, omega2, omega1h, omega2h, t, c, a0, b0
	print N, len(c)
	if N < len(c):
		reset_globals()
	for i in range(len(c) + 1, N+1):
		t[i] = (-1)^i*3*factorial(3*i-1)/(factorial(i)^3)
		c[i] = (3*i-1)*(3*i-2)*(3*i-3)*c[i-1] \
			+ (-1)^i*6*factorial(3*i-1)/(factorial(i)^3)*3*i^2 \
			- (-1)^i*6*factorial(3*i-4)/(factorial(i-1)^3)*\
				(81*(i-1)^2 + 54*(i-1) + 6)
		c[i] = -c[i]/i^3
		omega1h(z) = omega1h(z) + t[i]*z^i
		omega2h(z) = omega2h(z) + 2*ln(z)*t[i]*z^i + c[i]*z^i
	omega1(z) = omega1h(z)/(2*pi*I)
	omega2(z) = omega2h(z)/(4*pi^2*I^2)
	a0(z) = omega1(z) - 1/2
	b0(z) = -omega2(z)/2 + omega1(z)/2 - 1/4

def set_power_series_inf(N):
	global varpi1, varpi2, c1, c2, a_inf, b_inf
	print N, len(c1)
	if N < len(c1):
		reset_globals()
	omega = exp(2/3*pi*I)
	for i in range(len(c1) + 1, N+1):
		if (i % 3 == 0):
			c1[i] = 0
			c2[i] = 0
		else:
			c1[i] = (CC(3^i) * gamma(CC(i/3)))/\
				(gamma(CC(i+1)) * gamma(CC(1-i/3))^2)
			c2[i] = c1[i]*omega^i
		varpi1(psi) = varpi1(psi) + c1[i]*psi^i/(CC(2*pi*I))
		varpi2(psi) = varpi2(psi) + c2[i]*psi^i/(CC(2*pi*I))
	a_inf(psi) = varpi1(psi) - CC(1/2)
	b_inf(psi) = varpi1(psi)/3 - varpi2(psi)/3 - CC(1/3)

def test_inequality_ab(az, bz):
	a1 = float(real_part(az))
        a2 = float(imag_part(az))
        b1 = float(real_part(bz))
        b2 = float(imag_part(bz))
        B = - b2/a2
        d = -B*a1-b1+B^2/2
	#d = - b1
        print "a =", az, ", b=", bz
        print B, d
        return B, d


def test_inequality0(x, y):
	z = x + I*y
	az = CC(a0(z))
	bz = CC(b0(z))
	B, d = test_inequality_ab(az, bz)
	return B, d

def test_inequality_inf(x, y):
	p = CC(x, y)
	az = CC(a_inf(p))
	bz = CC(b_inf(p))
	return test_inequality_ab(az, bz)

# generate a random complex number z = x+I y with |z| < 1/27
def random_z_point():
	x = 1
	y = 1
	while x^2 + y^2 > 1/27^2:
		x = -1/27 + 2/27*random()
		y = -1/27 + 2/27*random()
	return x, y
		
		

# Generate _count_ random points in the region around z=0 (i.e.,
# with \abs{z} < 1/27) and test the inequality
def run_inequality_tests0(count, threshold):
	close = []
	dlist = []
	max = -100.0
	for i in range(count):
		x, y = random_z_point()
		d = test_inequality0(x, y)
		if d > threshold:
			print
			print "CLOSE:", x, y
			close.append([x,y])
			dlist.append(d)
		if d > max:
			max = d
	print "Close values: ", close
	print "Maximum: ", max
	return close, dlist


# generate a random complex number \psi = x+I y with |\psi| < 1
# and 2pi/3 < arg(\psi) < 4pi/3
def random_psi_point():
	a = 0
	x = 1
	y = 1
	while (a < 2*pi/3 ) and (a > -2*pi/3) or (x^2 + y^2 > 1):
		x = -1 + 2*random()
		y = -1 + 2*random()
		a = arg(CC(x,y))
	#print x, y
	return x, y

# generate a random complex number \psi = x+I y with |\psi| < 1
# and 2pi/3 < arg(\psi) < pi
def less_random_psi_point():
	a = 0
	x = 1
	y = 1
	while (a < 2*pi/3 ) or (x^2 + y^2 > 1):
		x = -1 + 2*random()
		y = -1 + 2*random()
		a = arg(CC(x,y))
	#print x, y
	return x, y

def run_simple_test_inf(count, threshold):
	close = []
	dlist = []
	maxD = -100.0
	for i in range(count):
		print i
		x, y = less_random_psi_point()
		p = CC(x, y)
		az = CC(a_inf(p))
		bz = CC(b_inf(p))
		a1 = float(real_part(az))
		a2 = float(imag_part(az))
		b1 = float(real_part(bz))
		b2 = float(imag_part(bz))
		B = - b2/a2
		if (B > 0.0) or (B < -0.5):
			print "Error!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
		print "psi = ", p, "Re a =", a1, "Re b = ", b1
		D = max([-b1, 0.5*a1-b1+1/8])
		print D
		if D > maxD:
			maxD = D
		if D > threshold:
			print "CLOSE:", x, y
			close.append([x,y])
			dlist.append(D)
	print "Maximum: ", maxD
	return close, dlist


# Run _count_ random test of the inequality in the psi-region
# Returns a list of all points with d > threshold, and
# a list with the corresponding values of d
def run_inequality_tests_inf(count, threshold):
	close = []
	dlist = []
	max = -100.0
	for i in range(count):
		print i
		x, y = random_psi_point()
		B, d = test_inequality_inf(x, y)
		if d > threshold:
			print
			print "CLOSE:", x, y
			close.append([x,y])
			dlist.append(d)
		if d > max:
			max = d
	#print "Close values: ", close
	print "Maximum: ", max
	return close, dlist
		

def find_stability_regions(count):
	points = {}
	for t1 in [True, False]:
		for t2 in [True, False]:
			for t3 in [True, False]:
				points[(t1, t2, t3)] = []
	for i in range(count):
		print i
		x = -1 + 2*random()
		y = -1 + 2*random()
		if (x^2 + y^2 < 1):
			psi = CC(x, y)
			O = b_inf(psi)
			OM1 = -2*b_inf(psi) + a_inf(psi) -1/2
			On1 = b_inf(psi) - a_inf(psi) - 1/2
			t1 = float(imag_part(O)) > 0
			t2 = float(imag_part(OM1)) > 0
			t3 = float(imag_part(On1)) > 0
			points[(t1, t2, t3)].append((x,y))
	return points

def draw_stability_regions(count):
	points = find_stability_regions(count)
	Q = list_plot((0,0))
	for t in points.keys():
		color = map(int, t)
		if points[t] != []:
			Q += list_plot(points[t], rgbcolor = color)
	Q.show()
	return Q, points


def psi_plane_in_theta1(count):
	theta_1 = []
	not_theta_1 = []
	for i in range(count):
		in_theta_1 = True
		x, y = random_psi_point()
		print x,y
		psi = CC(x, y)
		ZO = b_inf(psi)
		phi0 = arg(-ZO) + 1
		ZOmega = -2*b_inf(psi) + a_inf(psi) - 1/2
		phi1 = arg(-ZOmega) + 1
		ZOm1 = b_inf(psi) - a_inf(psi) - 1/2
		phi2 = arg(-ZOm1) + 1
		if abs(phi1-phi0) > 1:
			print "phi1-phi0 violation:", x, y
			in_theta_1 = False
		if abs(phi2-phi1) > 1:
			print "phi2-phi1 violation:", x, y
			in_theta_1 = False
		if abs(phi0-phi2) > 1:
			print "phi0-phi2 violation:", x, y
			in_theta_1 = False
		if in_theta_1:
			theta_1.append((x,y))
		else:
			not_theta_1.append((x,y))
	return theta_1, not_theta_1


def plot_d_inf(count):
	close_list, dlist = run_inequality_tests_inf(count, -100)
	return [(close_list[i][0], close_list[i][1], dlist[i])\
				for i in range(len(dlist))]

def plot_d_0(count):
	close_list, dlist = run_inequality_tests0(count, -100)
	return [(close_list[i][0], close_list[i][1], dlist[i])\
				for i in range(len(dlist))]
	

def plot_bz(count):
	b_list = []
	for i in range(count):
		x, y = random_psi_point()
		b = CC(b_inf(CC(x, y)))
		b_list.append((float(real_part(b)), float(imag_part(b))))
	for i in range(count):
		x, y = random_z_point()
		b = CC(b0(CC(x, y)))
		b_list.append((float(real_part(b)), float(imag_part(b))))
	return b_list

		

reset_globals()

N = 500
CC = ComplexField(100)

set_power_series0(N)
set_power_series_inf(N)