"""Extended functionality for [pyMOR](https://pymor.org/) parameters"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/api/pymor/parameters.ipynb.

# %% auto 0
__all__ = ['SimplifyFunctionalRules', 'simplify_functionals']

# %% ../../nbs/api/pymor/parameters.ipynb
from numbers import Number
from fastcore.basics import patch
import sympy as sy
from sympy.parsing.sympy_parser import parse_expr
from pymor.parameters.functionals import ProductParameterFunctional, ExpressionParameterFunctional, ParameterFunctional
from pymor.algorithms.rules import match_class, RuleTable, match_always
from pymor.basic import LincombOperator
from pymor.models.interface import Model
from pymor.operators.interface import Operator

# %% ../../nbs/api/pymor/parameters.ipynb
@patch
def __eq__(self:ExpressionParameterFunctional, other):
    return (
        isinstance(other, ExpressionParameterFunctional)
        and self.expression == other.expression
        and self.parameters == other.parameters
    )

# %% ../../nbs/api/pymor/parameters.ipynb
@patch
def __str__(self:ParameterFunctional):
    return f'f({", ".join(self.parameters)})'

# %% ../../nbs/api/pymor/parameters.ipynb
class SimplifyFunctionalRules(RuleTable):
    """|RuleTable| for the :func:`expand` algorithm."""

    def __init__(self):
        super().__init__(use_caching=True)

    @match_class(ProductParameterFunctional)
    def action_ProductParameterFunctional(self, functional):
        # merge child ProductParameterFunctional objects
        if any(isinstance(factor, ProductParameterFunctional) for factor in functional.factors):
            factors = []
            for factor in functional_factors:
                if isinstance(factor, ProductParameterFunctional):
                    factors.extend(self.apply(factor).factors)
                else:
                    factors.append(factor)
            functional = functional.with_(factors=factors)

        # multiply together numbers and ExpressionParameterFunctional objects
        if all(isinstance(factor, (ExpressionParameterFunctional, Number)) for factor in functional.factors):
            product = sy.prod([
                parse_expr(factor.expression) if isinstance(factor, ExpressionParameterFunctional) else factor 
                for factor in functional.factors
            ])
            if product.is_number:
                functional = product
            else:
                functional = ExpressionParameterFunctional(str(product), parameters=functional.parameters)
        
        return functional
    
    @match_class(LincombOperator)
    def action_LincombOperator(self, op):
        return op.with_(coefficients=[self.apply(c) for c in op.coefficients])

    @match_class(Model, Operator)
    def action_recurse(self, op):
        return self.replace_children(op)

    @match_always
    def action_generic(self, expr):
        return expr

# %% ../../nbs/api/pymor/parameters.ipynb
def simplify_functionals(obj):
    return SimplifyFunctionalRules().apply(obj)
