r'''
    Sage package for Functional order Basis.
'''
# Sage imports
from sage.all import cached_method, Integer, bessel_J, exp
from sage.all_cmdline import x

# Local imports
from ..psbasis import OrderBasis

class FunctionalBasis(OrderBasis):
    r'''
        Class for representing a basis of power functions.

        A basis of power functions is a type of order basis for power series
        where the `n`-th element is the `n`-th power of an order 1 power series `f(x)`.

        The first element in the sequence will always be the constant polynomial 1.

        The second element in the sequence is a function `f(x)` such that `f(0) = 0`
        and `f'(0) \neq 0`

        INPUT:
            - ``X``: the name for the operator representing the multiplication by `f(x)`.
    '''
    def __init__(self, f, X='f'):
        ## Initializing the PolyBasis structure
        super(FunctionalBasis,self).__init__()

        ## Adding the extra information
        self.__fun_name = X

        if(f(x=0) != 0):
            raise ValueError("The first initial value of the function has to be zero. Found %s" %f(x=0))
        self.__function = f

        ## The multiplication by X compatibility is given
        Sni = self.Sni()
        self.set_compatibility(X, Sni)

    def _element(self, n):
        r'''
            Method to return the `n`-th element of the basis.

            This method *implements* the corresponding abstract method from :class:`~pseries_basis.psbasis.PSBasis`.
            See method :func:`~pseries_basis.psbasis.PSBasis.element` for further information.

            This method removes the optional argument of but adds a new argument for getting 
            a symbolic representation as a polynomial in the function or the real object.

            For a :class:`FunctionalBasis` the output will be a function of order `n`.

            OUTPUT:

            A formal power series of order ``n``.

            TODO: FIX THIS METHOD
        '''
        return self.__function**n

    def __repr__(self):
        return "Functional Power Basis (%s)" %(self.__fun_name)

    def _latex_(self):
        return r"\left\{%s(x)^n\right\}_{n \geq 0}" %(self.__fun_name)

###############################################################
### EXAMPLES OF PARTICULAR GENERALIZED POWER BASIS
###############################################################
class ExponentialBasis(FunctionalBasis):
    r'''
        Class for the power functional basis generated by the exponential.

        This class represents the FunctionalBasis formed by the set of powers
        of the exponential function `f(x) = exp(x)-1`. The minus one only shows up
        because `f(x)` must have order 0.

        Following the notation in :arxiv:`2202.05550`, we can find that
        this basis has compatibilities with the multiplication by `exp(x)` and the
        derivation operator

        INPUT:
            - ``E``: the name for the operator representing the multiplication by `exp(x)`.
            - ``Dx``: the name for the operator representing the derivation by `x`.
    '''
    def __init__(self, E='E', Dx='Dx'):
        super(ExponentialBasis, self).__init__(exp(x)-1)

        Sni = self.Sni(); n = self.n(); Sn = self.Sn()

        self.set_compatibility(E, Sni + 1)
        self.set_derivation(Dx, n + (n+1)*Sn)



    def __repr__(self):
        return "Exponential basis (1, e^x-1, (e^x-1)^2,...)"

    def _latex_(self):
        return r"\left\{\left(e^x-1\right)^n\right\}_{n \geq 0}"

###############################################################
### EXAMPLES OF PARTICULAR ORDER BASIS
###############################################################
class BesselBasis(OrderBasis):
    r'''
        Class for representing a basis of Bessel functions.

        The Bessel functions with integer index form naturally an order basis. They are usually denoted
        by `J_n(x)` where `n` is a natural number.

        INPUT:
            - ``Xi``: the name for the operator representing the multiplication by `1/x`.
            - ``Dx``: the name of the operator representing the standard derivation.
    '''
    def __init__(self, Xi='Xi', Dx='Dx'):
        ## Initializing the PolyBasis structure
        super(BesselBasis,self).__init__()

        ## The multiplication by X compatibility is given
        Sni = self.Sni(); n = self.n(); Sn = self.Sn(); Q12 = 1/Integer(2)
        self.set_compatibility(Xi, (Q12/n)*Sn + (Q12/n)*Sni)
        self.set_derivation(Dx, Q12*Sn - Q12*Sni)

    @cached_method
    def element(self, n):
        r'''
            Method to return the `n`-th element of the basis.

            This method *implements* the corresponding abstract method from :class:`~pseries_basis.psbasis.PSBasis`.
            See method :func:`~pseries_basis.psbasis.PSBasis.element` for further information.

            This method removes the optional argument of variable name. 

            For a :class:`BesselBasis` the output will be the `n`-th Bessel function `J_n(x)`.

            OUTPUT:

            The `n`-th Bessel function that has order `n` as a formal power series.

            TODO: add examples
        '''
        if(n >= 0):
            return bessel_J(n,x)
        else:
            raise IndexError("The index must be a non-negative integer")

    def __repr__(self):
        return "Bessel Basis (J_n)"

    def _latex_(self):
        return r"\left\{J_n(x)\right\}_{n \geq 0}"

__all__ = ["FunctionalBasis", "ExponentialBasis", "BesselBasis"]