# ____________________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright (c) 2008-2026 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and Engineering
# Solutions of Sandia, LLC, the U.S. Government retains certain rights in this
# software.  This software is distributed under the 3-clause BSD License.
# ____________________________________________________________________________________

from pyomo.core.expr import identify_variables, value, differentiate
import logging
import math

logger = logging.getLogger(__name__)


def _loop(derivs, e_vars, diff_mode, ndx_list):
    for ndx, item in enumerate(derivs):
        ndx_list.append(ndx)
        if item.__class__ == list:
            for a, b in _loop(item, e_vars, diff_mode, ndx_list):
                yield a, b
        else:
            _derivs = differentiate(item, wrt_list=e_vars, mode=diff_mode)
            derivs[ndx] = _derivs
            yield ndx_list, _derivs
        ndx_list.pop()


def taylor_series_expansion(
    expr, diff_mode=differentiate.Modes.reverse_numeric, order=1
):
    """
    Generate a taylor series approximation for expr.

    Parameters
    ----------
    expr: pyomo.core.expr.numeric_expr.NumericExpression
    diff_mode: pyomo.core.expr.calculus.derivatives.Modes
        The method for differentiation.
    order: The order of the taylor series expansion
        If order is not 1, then symbolic differentiation must
        be used (differentiation.Modes.reverse_symbolic or
        differentiation.Modes.sympy).

    Returns
    -------
    res: pyomo.core.expr.numeric_expr.NumericExpression
    """
    if order < 0:
        raise ValueError(
            'Cannot compute taylor series expansion of order {0}'.format(str(order))
        )
    if order != 1 and diff_mode is differentiate.Modes.reverse_numeric:
        logger.warning(
            'taylor_series_expansion can only use symbolic differentiation for orders larger than 1'
        )
        diff_mode = differentiate.Modes.reverse_symbolic
    e_vars = list(identify_variables(expr=expr, include_fixed=False))

    res = value(expr)
    if order >= 1:
        derivs = differentiate(expr=expr, wrt_list=e_vars, mode=diff_mode)
        res += sum(
            (e_vars[i] - e_vars[i].value) * value(derivs[i]) for i in range(len(e_vars))
        )

    """
    This last bit of code is just for higher order taylor series expansions.
    The recursive function _loop modifies derivs in place so that derivs becomes a 
    list of lists of lists... However, _loop is also a generator so that 
    we don't have to loop through it twice. _loop yields two lists. The 
    first is a list of indices corresponding to the first k-1 variables that
    differentiation is being done with respect to. The second is a list of 
    derivatives. Each entry in this list is the derivative with respect to 
    the first k-1 variables and the kth variable, whose index matches the 
    index in _derivs.
    """
    if order >= 2:
        for n in range(2, order + 1):
            coef = 1.0 / math.factorial(n)
            for ndx_list, _derivs in _loop(derivs, e_vars, diff_mode, list()):
                tmp = coef
                for ndx in ndx_list:
                    tmp *= e_vars[ndx] - e_vars[ndx].value
                res += tmp * sum(
                    (e_vars[i] - e_vars[i].value) * value(_derivs[i])
                    for i in range(len(e_vars))
                )

    return res
