import pytest
import numpy as np
from scipy import cos, sin

from cxroots import Rectangle

@pytest.mark.parametrize('useDerivative', [True, False])
def test_count_roots(useDerivative):
	"""
	Example from "Locating all the Zeros of an Analytic Function in one Complex Variable"
	M.Dellnitz, O.Schutze, Q.Zheng, J. Compu. and App. Math. (2002), Vol.138, Issue 2

	There should be 424 roots inside this contour
	"""
	C  = Rectangle([-20.3,20.7], [-20.3,20.7])
	f  = lambda z: z**50 + z**12 - 5*sin(20*z)*cos(12*z) - 1
	df = lambda z: 50*z**49 + 12*z**11 + 60*sin(12*z)*sin(20*z) - 100*cos(12*z)*cos(20*z)

	if not useDerivative: df = None

	assert C.count_roots(f, df, verbose=True) == 424


@pytest.mark.xfail(reason='Possibly lack of high precision arithmetic')
def test_RingOscillator():
	# problem in section 4.3 of [DSZ]
	from cxroots import Rectangle, findRoots

	def A(z):
		t = 2.
		return np.array([[-0.0166689-2.12e-14*z, 1/60. + 6e-16*exp(-t*z)*z],
		 		  		 [0.0166659+exp(-t*z)*(-0.000037485+6e-16*z), -0.0166667-6e-16*z]])

	def dA(z):
		t = 2.
		return np.array([[-2.12e-14*np.ones_like(z), 6e-16*exp(-t*z) - t*6e-16*exp(-t*z)*z],
				   		 [-t*exp(-t*z)*(-0.000037485+6e-16*z)+exp(-t*z)*6e-16, -6e-16*np.ones_like(z)]])
	
	def f(z):
		AVal = np.rollaxis(A(z),-1,0)
		return np.linalg.det(AVal)

	def df(z):
		AVal  = A(z)
		dAVal = dA(z)
		return dAVal[0,0]*AVal[1,1] + AVal[0,0]*dAVal[1,1] - dAVal[0,1]*AVal[1,0] - AVal[0,1]*dAVal[1,0]


	box = Rectangle([-12,0], [-40,40])
	# roots_fdf = findRoots(box, f, df)
	# roots_f = findRoots(box, f)

	# XXX: No roots are recorded within the initial contour.
	# 	Perhaps because the coefficents of z are very small?
	# 	Perhaps need higher precision?
	assert box.count_enclosed_roots(f, df) != 0
	assert box.count_enclosed_roots(f) != 0

	# compare with fig 4 of [DSZ]
	# roots_fdf.show()
