from __future__ import division
import numpy as np
from numpy import inf, pi
import scipy.integrate
import scipy.misc
import warnings
import numdifftools.fornberg as ndf
import numdifftools

def prod(C, f, df=None, phi=None, psi=None, absTol=1e-12, relTol=1e-12, divMin=3, 
	divMax=15, m=2, intMethod='quad', integerTol=inf, verbose=False, callback=None):
	r"""
	Compute the symmetric bilinear form used in (1.12) of [KB]_.

	.. math::

		<\phi,\psi> = \frac{1}{2\pi i} \oint_C \phi(z) \psi(z) \frac{f'(z)}{f(z)} dz.
	
	Parameters
	----------
	C : :class:`Contour <cxroots.Contour.Contour>`
		A contour in the complex plane for.  No roots or poles of f
		should lie on C.
	f : function
		Function of a single variable f(x)
	df : function, optional
		Function of a single variable, df(x), providing the derivative 
		of the function f(x) at the point x.  If not provided then df is 
		approximated using a finite difference method.
	phi : function, optional
		Function of a single variable phi(x).  If not provided then 
		phi(z)=1.
	psi : function, optional
		Function of a single variable psi(x).  If not provided then 
		psi(z)=1.
	absTol : float, optional
		Absolute error tolerance for integration.
	relTol : float, optional
		Relative error tolerance for integration.
	divMin : int, optional
		Only used if intMethod='romb'. Minimum number of divisions before 
		the Romberg integration routine is allowed to exit.  
	divMax : int, optional
		Only used if intMethod='romb'.  The maximum number of divisions 
		before the Romberg integration routine of a path exits.  
	m : int, optional
		Only used if df=None and intMethod='quad'.  Must be even.  The 
		argument order=m is passed to numdifftools.Derivative and is the 
		order of the error term in the Taylor approximation.
	intMethod : {'quad', 'romb'}, optional
		If 'quad' then scipy.integrate.quad is used to perform the 
		integral.  If 'romb' then Romberg integraion, using 
		scipy.integrate.romb, is performed instead.
	integerTol : float, optional
		Only used when intMethod is 'romb'.  The integration routine will 
		not exit unless the result is within integerTol of an integer.  
		This is useful when computing the number of roots in a contour,
		which must be an integer.  By default integerTol is inf.
	verbose : bool, optional
		If True runtime information will be printed.  False be default.
	callback : function, optional
		Only used when intMethod is 'romb'.  A function that at each 
		step in the iteration is passed the current approximation for 
		the integral, the estimated error of that approximation and the
		number of iterations.  If the return of callback evaluates to 
		True then the integration will end. 

	Returns
	-------
	complex
		The value of the integral <phi, psi>.
	float
		An estimate of the error for the integration.

	References
	----------
	.. [KB] "Computing the zeros of analytic functions" by Peter Kravanja, 
		Marc Van Barel, Springer 2000
	"""
	if intMethod == 'romb':
		N = 1
		k = 0
		I = []

		while k < divMax and (len(I) < divMin
			or (abs(I[-2] - I[-1]) > absTol and abs(I[-2] - I[-1]) > relTol*abs(I[-1]))
			or (abs(I[-3] - I[-2]) > absTol and abs(I[-3] - I[-2]) > relTol*abs(I[-2]))
			or abs(int(round(I[-1].real)) - I[-1].real) > integerTol 
			or abs(I[-1].imag) > integerTol):
			N = 2*N
			t = np.linspace(0,1,N+1)
			k += 1
			dt = t[1]-t[0]

			integrals = []
			for segment in C.segments:
				# compute/retrieve function evaluations
				fVal = segment.trap_values(f,k)

				if df is None:
					# approximate df/dz with finite difference
					dfdt = np.gradient(fVal, dt)
					dfVal = dfdt/segment.dzdt(t)
				else:
					dfVal = segment.trap_values(df,k)

				segment_integrand = dfVal/fVal*segment.dzdt(t)
				if phi is not None:
					segment_integrand = segment.trap_values(phi,k)*segment_integrand
				if psi is not None:
					segment_integrand = segment.trap_values(psi,k)*segment_integrand

				segment_integral = scipy.integrate.romb(segment_integrand, dx=dt, axis=-1)/(2j*pi)
				integrals.append(segment_integral)
			
			I.append(sum(integrals))

			if verbose:
				if k > 1:
					print(k, 'I', I[-1], 'err', I[-2] - I[-1])
				else:
					print(k, 'I', I[-1])

			if callback is not None:
				err = abs(I[-2] - I[-1]) if k > 1 else None
				if callback(I[-1], err, k):
					break

		return I[-1], abs(I[-2] - I[-1])

	elif intMethod == 'quad':
		if df is None:
			df = numdifftools.Derivative(f, order=m)
			# df = lambda z: scipy.misc.derivative(f, z, dx=1e-8, n=1, order=3)
			
			### Too slow
			# ndf.derivative returns an array [f, f', f'', ...]
			# df = np.vectorize(lambda z: ndf.derivative(f, z, n=1)[1])

		I, err = 0, 0
		for segment in C.segments:
			integrand_cache = {}
			def integrand(t):
				if t in integrand_cache.keys():
					i = integrand_cache[t]
				else:
					z = segment(t)
					i = (df(z)/f(z))/(2j*pi) * segment.dzdt(t)
					if phi is not None:
						i = phi(z)*i
					if psi is not None:
						i = psi(z)*i
					integrand_cache[t] = i
				return i

			# integrate real part
			integrand_real = lambda t: np.real(integrand(t))
			result_real = scipy.integrate.quad(integrand_real, 0, 1, full_output=1, epsabs=absTol, epsrel=relTol)
			I_real, abserr_real, infodict_real = result_real[:3]

			# integrate imaginary part			
			integrand_imag = lambda t: np.imag(integrand(t))
			result_imag = scipy.integrate.quad(integrand_imag, 0, 1, full_output=1, epsabs=absTol, epsrel=relTol)
			I_imag, abserr_imag, infodict_imag = result_imag[:3]

			I   += I_real + 1j*I_imag
			err += abserr_real + 1j*abserr_imag

		return I, abs(err)

	else:
		raise ValueError("intMethod must be either 'romb' or 'quad'")

class RootError(RuntimeError):
	pass

def count_roots(C, f, df=None, NIntAbsTol=0.07, integerTol=0.1, divMin=3, 
	divMax=15, m=2, intMethod='quad', verbose=False):
	r"""
	For a function of one complex variable, f(z), which is analytic in 
	and within the contour C, return the number of zeros (counting 
	multiplicities) within the contour, N, using Cauchy's argument 
	principle,
	
	.. math::

		N = \frac{1}{2i\pi} \oint_C \frac{f'(z)}{f(z)} dz.

	If df(z), the derivative of f(z), is provided then the above 
	integral is computed directly.  Otherwise the derivative is 
	approximated using a finite difference method.

	The number of roots is taken to be the closest integer to the 
	computed value of the integral and the result is only accepted 
	if the integral is within integerTol of the closest integer.
	
	Parameters
	----------
	C : :class:`Contour <cxroots.Contour.Contour>`
		The contour which encloses the roots of f(z) that are to be 
		counted. 
	f : function
		Function of a single variable f(z).
	df : function, optional
		Function of a single complex variable, df(z), providing the 
		derivative of the function f(z) at the point z.  If not 
		provided, df will be approximated using a finite difference 
		method.
	NIntAbsTol : float, optional
		Required absolute error tolerance for the contour integration.
		Since the Cauchy integral must be an integer it is only 
		necessary to distinguish which integer the integral is 
		converging towards.  Therefore, NIntAbsTol can be fairly large.
	integerTol : float, optional
		The evaluation of the Cauchy integral will be accepted if its 
		value is within integerTol of the closest integer.
	divMin : int, optional
		Only used if intMethod='romb'. Minimum number of divisions 
		before the Romberg integration routine is allowed to exit.
	divMax : int, optional
		Only used if intMethod='romb'.  The maximum number of divisions 
		before the Romberg integration routine of a path exits.
	m : int, optional
		Only used if df=None and intMethod='quad'.  The argument order=m 
		is passed to numdifftools.Derivative and is the order of the 
		error term in the Taylor approximation.  m must be even.
	intMethod : {'quad', 'romb'}, optional
		If 'quad' then scipy.integrate.quad is used to perform the 
		integral.  If 'romb' then Romberg integraion, using 
		scipy.integrate.romb, is performed instead.
	verbose : bool, optional
		If True certain messages regarding the integration will be 
		printed.

	Returns
	-------
	int
		The number of zeros of f (counting multiplicities) which lie 
		within the contour C.
	"""
	if verbose:
		print('Computing number of roots within', C)

	with warnings.catch_warnings():
		# ignore warnings and catch if I is NaN later
		warnings.simplefilter("ignore")
		I, err = prod(C, f, df, absTol=NIntAbsTol, relTol=0, divMin=divMin, 
			divMax=divMax, m=m, intMethod=intMethod, verbose=verbose, integerTol=integerTol)

	if intMethod == 'romb':
		C._numberOfDivisionsForN = int(np.log2(len(C.segments[0]._trapValuesCache[f])-1))

	if np.isnan(I):
		raise RootError("""Result of integral is an invalid value.  
						   Most likely because of a divide by zero error.""")

	elif abs(int(round(I.real)) - I.real) < integerTol and abs(I.imag) < integerTol:
		# integral is sufficiently close to an integer
		numberOfZeros = int(round(I.real))
		return numberOfZeros

	else:
		raise RootError("The number of enclosed roots has not converged to an integer")

