from sympy import Symbol, symbols


def get_coordinates(function):
    """ Given a function evaluated from finite difference: return the coordinate system and its differential forms

    Example
    -------

        >>> t, x, y = symbols('t x y')
        >>> dt, dx, dy = symbols('dt dx dy')
        >>> U = Function('U')(t,x,y)
        >>> U.subs({x:x+dx, y:y-3dy/2})
        U(t,x+dx,y-3dy/2)
        >>> get_coordinates(U)
        ([t,x,y],[dt,dx,dy])

    """

    x = []
    dx = []

    for arg in function.args:

        coordinate = list(arg.atoms(Symbol))

        if len(coordinate) == 2:
            xi, dxi = coordinate
            if str(xi)[0] == 'd':
                xi, dxi = dxi, xi
        elif len(coordinate) == 1:
            xi = coordinate.pop()
            dxi = symbols('d' + str(xi))
        else:
            raise ValueError(f"Too many symbols {arg} for function {function}")

        x.append(xi)
        dx.append(dxi)

    return x, dx