
##############################################################################
#
# Copyright (c) 2003-2026 by the esys.escript Group
# https://github.com/LutzGross/esys-escript.github.io
#
# Primary Business: Queensland, Australia
# Licensed under the Apache License, version 2.0
# http://www.apache.org/licenses/LICENSE-2.0
#
# See CREDITS file for contributors and development history
#
##############################################################################

"""
Utility functions for symbolic computations in escript.

This module provides helper functions for working with symbolic expressions,
including creating Symbol objects and computing total differentials.

:var __author__: name of author
:var __copyright__: copyrights
:var __license__: licence agreement
:var __url__: url entry point on documentation
:var __version__: version
:var __date__: date of the version
"""


__copyright__="""Copyright (c) 2003-2026 by the esys.escript Group
https://github.com/LutzGross/esys-escript.github.io
Primary Business: Queensland, Australia"""
__license__="""Licensed under the Apache License, version 2.0
http://www.apache.org/licenses/LICENSE-2.0"""
__url__="https://github.com/LutzGross/esys-escript.github.io"
__author__="Cihan Altinay"

import numpy
import sympy
from .symbol import Symbol

def symbols(*names, **kwargs):
    """
    Creates one or more Symbol objects, emulating sympy.symbols behavior.

    :param names: symbol name(s), space or comma separated
    :type names: ``str``
    :keyword shape: shape of the symbol (default: scalar)
    :type shape: ``tuple`` of ``int``
    :keyword dim: spatial dimension
    :type dim: ``int``
    :return: single Symbol or tuple of Symbols
    :rtype: `Symbol` or ``tuple`` of `Symbol`
    """

    shape=kwargs.pop('shape', ())

    s = names[0]
    if not isinstance(s, list):
        import re
        s = re.split(r'\s|,', s)
    res = []
    for t in s:
        # skip empty strings
        if not t:
            continue
        sym = Symbol(t, shape, **kwargs)
        res.append(sym)
    res = tuple(res)
    if len(res) == 0:   # var('')
        res = None
    elif len(res) == 1: # var('x')
        res = res[0]
                        # otherwise var('a b ...')
    return res


def isSymbol(arg):
    """
    Tests if the argument is a symbolic object.

    :param arg: object to test
    :type arg: any
    :return: ``True`` if ``arg`` is an escript `Symbol` or ``sympy.Basic``
             object, ``False`` otherwise
    :rtype: ``bool``
    """
    return isinstance(arg, Symbol) or isinstance(arg, sympy.Basic)

def removeFsFromGrad(sym):
    """
    Removes function space parameters from gradient expressions.

    Returns ``sym`` with all occurrences of grad_n(a,b,c) replaced by
    grad_n(a,b). This strips the function space parameter from gradient calls.

    :param sym: symbolic expression to process
    :type sym: `Symbol`
    :return: expression with function space parameters removed from gradients
    :rtype: `Symbol`
    """
    from esys.escript import symfn
    gg=sym.atoms(symfn.grad_n)
    for g in gg:
        if len(g.args)==3:
            r=symfn.grad_n(*g.args[:2])
            sym=sym.subs(g, r)
    return sym

def getTotalDifferential(f, x, order=0):
    """
    Computes the total differential of f with respect to x.

    This function computes::

        | Df/Dx = del_f/del_x + del_f/del_grad(x)*del_grad(x)/del_x + ...
        |            \\   /         \\   /
        |              a             b

    :param f: the function to differentiate
    :type f: `Symbol` or ``numpy.ndarray``
    :param x: the variable to differentiate with respect to
    :type x: `Symbol`
    :param order: order of differentiation (0 or 1)
    :type order: ``int``
    :return: tuple of derivative terms (a,) for order=0, (a, b) for order=1
    :rtype: `Symbol` or ``tuple`` of `Symbol`
    """

    from .. import util
    res=()
    shape=util.getShape(f)
    if not isSymbol(f):
        res+=(numpy.zeros(shape+x.getShape()),)
        for i in range(order):
            x=x.grad()
            res+=numpy.zeros(shape+x.getShape())

    elif x.getRank()==0:
        f=removeFsFromGrad(f)
        dfdx=f.diff(x)
        dgdx=x.grad().diff(x)
        a=numpy.empty(shape, dtype=object)
        if order>0:
            b=numpy.empty(shape+dgdx.getShape(), dtype=object)

        if len(shape)==0:
            for j in numpy.ndindex(dgdx.getShape()):
                y=dfdx
                z=dgdx[j]
                # expand() and coeff() are very expensive so
                # we set the unwanted factors to zero to extract
                # the one we need
                for jj in numpy.ndindex(dgdx.getShape()):
                    if j==jj: continue
                    y=y.subs(dgdx[jj], 0)
                a=y.subs(z,0) # terms in x and constants
                if order>0:
                    b[j]=y.subs(z,1)-a
        else:
            for i in numpy.ndindex(shape):
                for j in numpy.ndindex(dgdx.getShape()):
                    y=dfdx[i]
                    z=dgdx[j]
                    for jj in numpy.ndindex(dgdx.getShape()):
                        if j==jj: continue
                        y=y.subs(dgdx[jj], 0)
                    a[i]=y.subs(z,0) # terms in x and constants
                    if order>0:
                        b[i+j]=y.subs(z,1)-a[i]
        res+=(Symbol(a, dim=f.getDim(), subs=f.getDataSubstitutions()),)
        if order>0:
            res+=(Symbol(b, dim=f.getDim(), subs=f.getDataSubstitutions()),)

    elif x.getRank()==1:
        f=removeFsFromGrad(f)
        dfdx=f.diff(x)
        dgdx=x.grad().diff(x).transpose(2)
        a=numpy.empty(shape+x.getShape(), dtype=object)
        if order>0:
            b=numpy.empty(shape+x.grad().getShape(), dtype=object)

        if len(shape)==0:
            raise NotImplementedError('f scalar, x vector')
        else:
            for i in numpy.ndindex(shape):
                for k,l in numpy.ndindex(x.grad().getShape()):
                    if dgdx[k,k,l]==0:
                        a[i+(k,)]=0
                        if order>0:
                            b[i+(k,l)]=0
                    else:
                        y=dfdx[i+(k,)]
                        z=dgdx[k,k,l]
                        for kk,ll in numpy.ndindex(x.grad().getShape()):
                            if k==kk and l==ll: continue
                            y=y.subs(dgdx[kk,kk,ll], 0)
                        a[i+(k,)]=y.subs(z,0) # terms in x and constants
                        if order>0:
                            b[i+(k,l)]=y.subs(z,1)-a[i+(k,)]

        res+=(Symbol(a, dim=f.getDim(), subs=f.getDataSubstitutions()),)
        if order>0:
            res+=(Symbol(b, dim=f.getDim(), subs=f.getDataSubstitutions()),)

    if len(res)==1:
        return res[0]
    else:
        return res
