from sage.all import *

class SymbolSequence:
    """
    A callable object which imitates a function from ZZ^n to a
    sequence with n subscripts.

    INPUT:

      - ``name`` -- The sequence name.

      - ``latex_name`` -- An optional latex expression (string) to
        use instead of `name` when converting the symbols to latex.

      - ``domain`` -- A string representing the domain of the symbol,
        either 'real', 'complex', or 'positive'.

    OUTPUT:

    A callable object returning symbolic expressions.

    EXAMPLES:

    Create coefficients for polynomials of arbitrary degree::

        sage: a = SymbolSequence('a')
        sage: p = sum([ a(i)*x^i for i in range(0,5)])
        sage: p
        a4*x^4 + a3*x^3 + a2*x^2 + a1*x + a0

    Using a different latex name since 'lambda' is reserved::

        sage: l = SymbolSequence('l', '\lambda')
        sage: l(0)
        l0
        sage: latex(l(0))
        \lambda_{0}

    Using multiple indices::

        sage: a = SymbolSequence('a')
        sage: a(0,1,2)
        a012
        sage: latex(a(0,1,2))
        a_{0}_{1}_{2}
        sage: [ a(i,j) for i in range(0,2) for j in range(0,2) ]
        [a00, a01, a10, a11]

    If no index is given, an unsubscripted symbol is returned::

        sage: a = SymbolSequence('a')
        sage: a()
        a

    You can pass slice objects instead of integers to obtain a list of
    symbols::

        sage: a = SymbolSequence('a')
        sage: a(slice(5,7))
        [a5, a6]

    This even works for the second, third, etc. indices::

        sage: a = SymbolSequence('a')
        sage: a(slice(0,2), slice(0,2))
        [a00, a01, a10, a11]

    You can also index with the list index operator::

        sage: a = SymbolSequence('a')
        sage: a[1]
        a1

    This allows you to retrieve one-dimensional slices easily::

        sage: a = SymbolSequence('a')
        sage: a[0:5]
        [a0, a1, a2, a3, a4]

    TESTS:

    We shouldn't overwrite variables in the global namespace::

        sage: a = SymbolSequence('a')
        sage: a0 = 4
        sage: a(0)
        a0
        sage: a0
        4

    The symbol at a given index should always be the same, even when
    the symbols themselves are unnamed. We store the string
    representation and compare because the output is unpredictable::

        sage: a = SymbolSequence()
        sage: a0str = str(a(0))
        sage: str(a(0)) == a0str
        True

    Slices and single indices work when combined::

        sage: a = SymbolSequence('a')
        sage: a(3, slice(0,2))
        [a30, a31]
        sage: a(slice(0,2), 3)
        [a03, a13]

    """

    def __init__(self, name=None, latex_name=None, domain=None):
        # We store a dict of already-created symbols so that we don't
        # recreate a symbol which already exists. This is especially
        # helpful when using unnamed variables, if you want e.g. a(0)
        # to return the same variable each time.
        #
        # The entry corresponding to None is the un-subscripted symbol
        # with our name.
        unsubscripted = SR.symbol(name, latex_name, domain)
        self._symbols = { None: unsubscripted }

        self._name = name
        self._latex_name = latex_name
        self._domain = domain


    def _create_symbol_(self, subscript):
        if self._name is None:
            # Allow creating unnamed symbols, for consistency with
            # SR.symbol().
            name = None
        else:
            name = '%s%d' % (self._name, subscript)

        if self._latex_name is None:
            latex_name = None
        else:
            latex_name = r'%s_{%d}' % (self._latex_name, subscript)

        return SR.symbol(name, latex_name, self._domain)


    def _flatten_list_(self, l):
        """
        Recursively flatten the given list, allowing for some elements
        to be non-iterable.
        """
        result = []

        for item in l:
            if isinstance(item, list):
                result += self._flatten_list_(item)
            else:
                result += [item]

        return result


    def __getitem__(self, key):
        if isinstance(key, slice):
            # We were given a slice. Clean up some of its properties
            # first. The start/step are default for lists. We make
            # copies of these because they're read-only.
            (start, step) = (key.start, key.step)
            if start is None:
                start = 0
            if key.stop is None:
                # Would otherwise loop forever since our "length" is
                # undefined.
                raise ValueError('You must supply an terminal index')
            if step is None:
               step = 1

            # If the user asks for a slice, we'll be returning a list
            # of symbols.
            return [ self(idx) for idx in range(start, key.stop, step) ]
        else:
            return self(key)


    def __call__(self, *args):
        args = list(args)

        if len(args) == 0:
            return self._symbols[None]

        # This is safe after the len == 0 test.
        key = args[0]
        args = args[1:] # Peel off the first arg, which we've called 'key'

        if isinstance(key, slice):
            if len(args) == 0:
                return self[key]
            else:
                v = self[key]
                ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
                       for w in v ]

                # This might be nested...
                maybe_nested_list = [ s(*args) for s in ss ]
                return self._flatten_list_(maybe_nested_list)

        if key < 0:
            # Cowardly refuse to create a variable named "a-1".
            raise IndexError('Indices must be nonnegative')

        if len(args) == 0:
            # Base case, create a symbol and return it.
            try:
                return self._symbols[key]
            except KeyError:
                self._symbols[key] = self._create_symbol_(key)
                return self._symbols[key]
        else:
            # If we're given more than one index, we want to create the
            # subsequences recursively. For example, if we're asked for
            # x(1,2), this is really SymbolSequence('x1')(2).
            v = self(key) # x(1) -> x1
            ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
            return ss(*args)