
import numpy as np
from rkstiff.solver import StiffSolverAS,StiffSolverCS

def phi1(z):
    """ Computes RKETD psi-function of the first order. 
    INPUTS
        z - real or complex-valued input array 
    OUTPUT
        return (exp(z)-1)/z  -> real or complex-valued RKETD function of first order
    """
    return (np.exp(z)-1)/z

def phi2(z):
    """ Computes RKETD psi-function of the second order. 
    INPUTS
        z - real or complex-valued input array 
    OUTPUT
        return 2!*(exp(z)-1-z/2)/z^2  -> real or complex-valued RKETD function of second order
    """
    return 2*(np.exp(z)-1-z)/z**2

def phi3(z):
    """ Computes RKETD psi-function of the third order. 
    INPUTS
        z - real or complex-valued input array 
    OUTPUT
        return 3!*(exp(z)-1-z/2-z^3/6)/z^3  -> real or complex-valued RKETD function of third order
    """
    return 6*(np.exp(z)-1-z-z**2/2)/z**3

class ETDAS(StiffSolverAS):
    """
    Class template for an RKETD adaptive-step solver. Adds several ETD specific parameters
    to the StiffSolverAS class.

    ATTRIBUTES
    ----------

    modecutoff : float
       psi RKETD functions are numerically inaccurate near zero when computed directly; when |z| < modecutoff use a
       contour integral (more computationally expensive method) to compute these values
    contour_points : int
        number of points to use in contour integral when evaluating small input values to the psi RKETD functions
    contour_radius : float
        radius size of a circular contour integral to be used to evaluate small input values to the psi RKETD functions

    """

    def __init__(self,linop,NLfunc,epsilon=1e-4,incrF = 1.25, decrF = 0.85, safetyF = 0.8,\
            adapt_cutoff = 0.01, minh = 1e-16, modecutoff = 0.01, contour_points = 32,\
            contour_radius = 1.0):
        super().__init__(linop,NLfunc,epsilon=epsilon,incrF=incrF,decrF=decrF,\
                safetyF=safetyF,adapt_cutoff=adapt_cutoff,minh=minh)
        self.modecutoff = modecutoff
        if (self.modecutoff > 1.0) or (self.modecutoff <= 0):
            raise ValueError('modecutoff must be between 0.0 and 1.0 but is {}'.format(self.modecutoff))
        self.contour_points = contour_points
        if not isinstance(self.contour_points,int):
            raise TypeError('contour_points must be an integer but is {}'.format(self.contour_points))
        if self.contour_points <= 1:
            raise ValueError('contour_points must be an integer greater than 1 but is {}'.format(self.contour_points))

        self.contour_radius = contour_radius
        if self.contour_radius <= 0:
            raise ValueError('contour_radius must greater than 0 but is {}'.format(self.contour_radius))
        self._h_coeff = None
        

class ETDCS(StiffSolverCS):
    """
    Class template for an RKETD constant-step solver. Adds several ETD specific parameters
    to the StiffSolverCS class.

    ATTRIBUTES
    ----------

    modecutoff : float
       psi RKETD functions are numerically inaccurate near zero when computed directly; when |z| < modecutoff use a
       contour integral (more computationally expensive method) to compute these values
    contour_points : int
        number of points to use in contour integral when evaluating small input values to the psi RKETD functions
    contour_radius : float
        radius size of a circular contour integral to be used to evaluate small input values to the psi RKETD functions

    """

    def __init__(self,linop,NLfunc,modecutoff = 0.01,contour_points = 32,contour_radius = 1.0):
        super().__init__(linop,NLfunc)
        self.modecutoff = modecutoff
        if (self.modecutoff > 1.0) or (self.modecutoff <= 0):
            raise ValueError('modecutoff must be between 0.0 and 1.0 but is {}'.format(self.modecutoff))
        self.contour_points = contour_points
        if not isinstance(self.contour_points,int):
            raise TypeError('contour_points must be an integer but is {}'.format(self.contour_points))
        if self.contour_points <= 1:
            raise ValueError('contour_points must be an integer greater than 1 but is {}'.format(self.contour_points))
        self.contour_radius = contour_radius
        if self.contour_radius <= 0:
            raise ValueError('contour_radius must greater than 0 but is {}'.format(self.contour_radius))
        self._h_coeff = None




