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

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

# %% ../../nbs/api/pymor/operators.ipynb 4
from __future__ import annotations
from pathlib import Path
import importlib
from numbers import Number
import math
from functools import reduce

from fastcore.basics import patch
from fastcore.meta import delegates
from pydash import unzip

import numpy as np
from numpy import array, ndarray
from scipy.sparse import issparse, sparray
from sparse import SparseArray
import sympy as sy
import pandas as pd
from xarray import DataArray, Coordinates

from pymor.basic import LincombOperator, NumpyMatrixOperator, ZeroOperator, ExpressionParameterFunctional, NumpyVectorSpace, Mu, IdentityOperator, Parameters
from pymor.models.interface import Model
from pymor.algorithms.to_matrix import to_matrix, ToMatrixRules
from pymor.algorithms.simplify import ExpandRules, expand, ContractRules, contract
from pymor.algorithms.rules import match_class, RuleTable
from pymor.algorithms.lincomb import assemble_lincomb
from pymor.operators.interface import Operator
from pymor.core.base import abstractmethod
from pymor.parameters.functionals import ParameterFunctional

from ..utilities.xarray import *
from ..utilities.sparse import sparse, sparse2d, kron, sparse2d_rand
from .parameters import *
from .vectorarrays import *

# %% auto 0
__all__ = ['operator_html_tree', 'coordinates_to_list', 'XarrayMatrixOperator', 'contract_matvec', 'dims_matvec', 'replace_item',
           'densify', 'XarrayFunctionalOperator', 'SumOperator', 'ScaleOperator', 'ProductOperator']

# %% ../../nbs/api/pymor/operators.ipynb
LincombOperator.suboperator_symbol = '+'

# %% ../../nbs/api/pymor/operators.ipynb
def _mapping_str(op:Operator):
    try:
        return f"{op.source._repr_html_(show_details=False)} → {op.range._repr_html_(show_details=False)}"
    except AttributeError:
        return f"{str(op.source)} → {str(op.range)}"

# %% ../../nbs/api/pymor/operators.ipynb
def _operator_group_tag(
    level:int, 
    cls:str, 
    mapping:str, 
    name:str|None=None,
    parameters:str|None=None,
    contents:str|None=None,
    suboperators:list(str)|None=None, 
    symbol:str|None=None, 
    opened:bool=False
) -> str:
    if suboperators or contents:
        return f'''
            <li>
                <details {'open' if opened else 'closed'}>
                    <summary>
                        <ul class="operator">
                            {'<li>' + name + '</li>' if name else ''}
                            <li>{cls}</li>
                            <li>{mapping}</li>
                            {'<li>' + parameters + '</li>' if parameters else ''}
                        </ul>
                    </summary>
                    <ul class="operator-group operator-group-{'even' if not (level % 2) else 'odd'}">
                        {contents or symbol.join(suboperators)}
                    </ul>
                </details>
            </li>
        '''
    else:
        return f'''
            <li>
                <ul class="operator">
                    {'<li>' + name + '</li>' if name else ''}
                    <li>{cls}</li>
                    <li>{mapping}</li>
                    {'<li>' + parameters + '</li>' if parameters else ''}
                </ul>
            </li>
        '''

# %% ../../nbs/api/pymor/operators.ipynb
def _operator_html_node(op:Operator, level:int=0, opened:bool=False) -> str:
    cls = op.__class__.__name__
    name = op.name if op.name != cls else None
    mapping = _mapping_str(op)
    parameters = str(list(Parameters.of(op)))
    contents = None
    if hasattr(op, 'operators'):
        suboperators = [_operator_html_node(o, level=level + 1, opened=opened) for o in op.operators]
        symbol = op.suboperator_symbol
        return _operator_group_tag(level, cls, mapping, name, parameters, contents, suboperators, symbol, opened)
    if hasattr(op, 'matrix'):
        contents = op.matrix._repr_html_() if hasattr(op.matrix, '_repr_html_') else str(op.matrix)
        return _operator_group_tag(level, cls, mapping, name, parameters, contents)
    return _operator_group_tag(level, cls, mapping, name, parameters)

# %% ../../nbs/api/pymor/operators.ipynb
def operator_html_tree(op:Operator, opened:bool=False) -> str:
    style = """
    <style>
        .operator-group summary {
            list-style: none;
            display: flex;
            gap: 0.25em;
            align-items: start;
            cursor: pointer;
        }
        .operator-group summary::before {
            content: '▶︎';
            font-size: 75%;
            align-items: center;
            width: 1em;
            flex-shrink: 0;
            transition: 0.3s;
        }
        .operator-group details[open] > summary::before {
            transform: rotate(90deg);
            transform-origin: 40% 45% 0;
        }
        .operator-group-odd > li {
            background-color: #f0f0f0;;
        }
        .operator-group-even > li {
            background-color: #fbfbfb;
        }
        .operator-group {
            display: flex;
            gap: .4em;
            align-items: center;
            padding: 0px;
            margin-left: -2em;
            margin-top: .25em;
        }
        .operator-group > li {
            list-style-type: none;
            border: 1px solid black;
            border-radius: .25em;
            padding: .25em;
        }
        .operator {
            margin: 0;
            padding: 0;
            margin-top: -.2em;
            margin-left: -2em;
        }
        .operator > li {
            list-style-type: none;
            border: none;
            padding: 0;
            line-height: 1.25em;
        }
    </style>
    """
    return style + f'<ul class="operator-group operator-group-odd">{_operator_html_node(op, opened=opened)}</ul>'

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def _repr_html_(self:Operator):
    return operator_html_tree(self)

# %% ../../nbs/api/pymor/operators.ipynb
class XarrayMatrixBasedOperator(Operator):
    """Base class for operators that assemble into an `XarrayMatrixOperator`."""

    linear = True

    @property
    def H(self):
        if not self.parametric:
            return self.assemble().H
        else:
            return super().H

    @abstractmethod
    def assemble(self, mu=None): pass

    def apply(self, U, mu=None):
        return self.range.from_xarray(self.assemble(mu).matrix.dot(U.impl._array))

    def apply_adjoint(self, V, mu=None):
        return self.assemble(mu).apply_adjoint(V)

    def as_range_array(self, mu=None):
        matrix = unlabel_dims(self.assemble(mu).matrix.squeeze())
        extended_coords = {d: matrix.coords[d] for d in matrix.dims if d not in self.range.dims}
        return self.range.from_numpy(self.matrix.data, extended_coords)

    def as_source_array(self, mu=None):
        matrix = unlabel_dims(self.assemble(mu).matrix.squeeze())
        extended_coords = {d: matrix.coords[d] for d in matrix.dims if d not in self.source.dims}
        return self.source.from_numpy(self.matrix.data, extended_coords).conj()

    @property
    def display_name(self):
        return self.name if self.name != self.__class__.__name__ else 'X'

    def __str__(self):
        return f"{self.display_name}{_mapping_str(self)}"

# %% ../../nbs/api/pymor/operators.ipynb
def _ndims(matrix):
    if isinstance(matrix, DataArray):
        return len(matrix.dims)
    else:
        return matrix.ndim    

# %% ../../nbs/api/pymor/operators.ipynb
def _label_dims_with(obj, lbl):
    if isinstance(obj, XarrayVectorSpace):
        return obj.with_(coord_data = _label_dims_with(obj.coord_data, lbl))
    if isinstance(obj, list | tuple): 
        return [d + lbl for d in obj]
    if isinstance(obj, dict):
        return {k + lbl: v for k,v in obj.items()}

# %% ../../nbs/api/pymor/operators.ipynb
def _label_dims(obj, src_dims, rng_dims):
    if isinstance(obj, XarrayVectorSpace):
        coord_data = {_label_dims(k, src_dims, rng_dims): v for k,v in obj.coord_data.items()}
        return obj.with_(coord_data=coord_data)
    if isinstance(src_dims, str): src_dims = [src_dims]
    if isinstance(rng_dims, str): rng_dims = [rng_dims]
    labled_src_dims = _label_dims_with(src_dims, Lbl.SRC)
    labled_rng_dims = _label_dims_with(rng_dims, Lbl.RNG)
    rename_dict = dict(zip(src_dims, labled_src_dims)) | dict(zip(rng_dims, labled_rng_dims))
    match obj:
        case str():
            return rename_dict.get(obj, obj)
        case list() | tuple():
            return [rename_dict.get(dim, dim) for dim in obj]
        case DataArray():
            return obj.rename(rename_dict)
        case Coordinates():
            return Coordinates({rename_dict.get(k, k): v.data for k,v in obj.items()})

# %% ../../nbs/api/pymor/operators.ipynb
def _swap_dim_labels(da):
    src_dims = [d for d in da.dims if d.endswith(Lbl.SRC)]
    rng_dims = [d for d in da.dims if d.endswith(Lbl.RNG)]
    rename_dict = {d: d.replace(Lbl.SRC, Lbl.RNG) for d in src_dims} | {d: d.replace(Lbl.RNG, Lbl.SRC) for d in rng_dims}
    return da.rename(rename_dict)

# %% ../../nbs/api/pymor/operators.ipynb
def _is_labeled_da(da):
    if not isinstance(da, DataArray): return False
    return any(d.endswith(Lbl.SRC) or d.endswith(Lbl.RNG) for d in da.dims)

# %% ../../nbs/api/pymor/operators.ipynb
def coordinates_to_list(coords, dims):
    return [(d, coords[d].data, coords[d].attrs) for d in dims]

# %% ../../nbs/api/pymor/operators.ipynb
def _get_space_dims(dims):
    if dims is None: return None, None
    if isinstance(dims, dict): dims = Coordinates(dims)
    if isinstance(dims, Coordinates): return list(dims), dims
    if isinstance(dims, str): dims = [dims]
    if isinstance(dims, list | tuple): return dims, None
    if hasattr(dims, 'dims'): return dims.dims, dims.coords
    raise ValueError("Source or range specification not valid")

# %% ../../nbs/api/pymor/operators.ipynb
class XarrayMatrixOperator(XarrayMatrixBasedOperator):
    """An `Operator` backed by an xarray `DataArray`."""

    def __init__(
        self, 
        matrix:DataArray|ndarray|SparseArray|sparray, # N-dimensional matrix
        source:XarrayVectorSpace|DataArray|Coordinates|dict|list|str|None=None, # Source vector space or coordinates
        range:XarrayVectorSpace|DataArray|Coordinates|dict|list|str|None=None, # Range vector space or coordinates
        solver_options:dict=None, # Options for matrix solver
        name:str=None # Operator name
    ):
        if issparse(matrix): matrix = sparse(matrix)
        if not isinstance(matrix, DataArray):
            if not isinstance(source, XarrayVectorSpace): 
                if isinstance(source, str): source = [source]
                if source is not None: 
                    source = validate_coord_data(source, matrix.shape[-len(source):])
                source = XarrayVectorSpace(source)
            if not isinstance(range, XarrayVectorSpace):
                if isinstance(range, str): range = [range]
                if range is not None:
                    range = validate_coord_data(range, matrix.shape[:len(range)])
                range = XarrayVectorSpace(range)
            labeled_source = _label_dims_with(source, Lbl.SRC)
            labeled_range = _label_dims_with(range, Lbl.RNG)
            matrix = DataArray(
                matrix, 
                dims=labeled_range.dims + labeled_source.dims,
                coords=make_coordinates(labeled_range.coord_data | labeled_source.coord_data),
                name=name
            )
        else:
            if not _is_labeled_da(matrix): matrix = _label_dims(matrix, source, range)
            coords = validate_coord_data(matrix)
            labeled_source = XarrayVectorSpace({k:v for k,v in coords.items() if k.endswith(Lbl.SRC)})
            labeled_range = XarrayVectorSpace({k:v for k,v in coords.items() if k.endswith(Lbl.RNG)})
            source = unlabel_dims(labeled_source)
            range = unlabel_dims(labeled_range)
        if name is None: name = matrix.name
        self.__auto_init(locals())
        self.labeled_source = labeled_source
        self.labeled_range = labeled_range

    def assemble(self, mu=None): return self

    @property
    def size(self):
        return self.matrix.size
        
    @property
    def sparse(self):
        return isinstance(self.matrix.data, SparseArray)

# %% ../../nbs/api/pymor/operators.ipynb
@patch(as_prop=True)
def H(self:XarrayMatrixOperator):
    """The adjoint operator."""
    adjoint_matrix = self.matrix if np.isrealobj(self.matrix) else self.matrix.conj()           
    options = {'inverse': self.solver_options.get('inverse_adjoint'),
               'inverse_adjoint': self.solver_options.get('inverse')} if self.solver_options else None
    return self.with_(matrix=_swap_dim_labels(adjoint_matrix), solver_options=options, name=self.name + '_adjoint')

# %% ../../nbs/api/pymor/operators.ipynb
def contract_matvec(a, b, axes):
    """Contract one axis of the first (2D) array with one axis of the second (ND) array."""
    out = np.tensordot(a, b, axes)
    if a.ndim == 1: return out
    return np.moveaxis(out, 0, axes[1])

# %% ../../nbs/api/pymor/operators.ipynb
def dims_matvec(a, b):
    """Contract one axis of the first (2D) array with the axis of the second (ND) array of the same name."""
    a, adims = a
    b, bdims = b
    ax = [[adims.index(d), i] for i,d in enumerate(bdims) if d in adims][0]
    newdims = list(bdims)
    if len(adims) == 1:
        newdims.pop(ax[1])
    else:
        newdims[ax[1]] = adims[1 - ax[0]]
    return contract_matvec(a, b, ax), newdims

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def __truediv__(self:XarrayVectorSpace, other):
    if not isinstance(other, XarrayVectorSpace): raise NotImplementedError
    return self.with_(coords={k: v for k,v in self.coords.items() if k not in other.coords})

# %% ../../nbs/api/pymor/operators.ipynb
def replace_item(dct, old_key, new_pair):
    return dict(new_pair if k == old_key else (k, v) for k,v in dct.items())

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def range_under(self:XarrayVectorSpace, op:XarrayMatrixOperator):
    """Find the space that this space is mapped to by an XarrayMatrixOperator. Assumes that the operator has one-dimensional source and range."""
    source_dim = op.source.dims[0]
    if not op.range.dims:
        coords = {k:v for k,v in self.coord_data.items() if k != source_dim}
    else:
        range_dim = op.range.dims[0]
        range_coord = op.range.coord_data[range_dim]
        coords = replace_item(self.coord_data, source_dim, (range_dim, range_coord))
    return self.with_(coord_data=coords)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def apply(
    self:XarrayMatrixOperator, 
    U, # `XarrayVectorArray` of vectors to which the operator is applied
    mu=None, # The parameter values for which to evaluate the operator
    method='numpy',
)->XarrayVectorArray: # `XarrayVectorArray` in the range `XarrayVectorSpace`
    """Apply the operator to an `XarrayVectorArray` in the source vector space."""
    assert U in self.source
    if method == 'numpy':
        src_dims = _label_dims(U.array.dims, U.dims, [])
        out, out_dims = dims_matvec((self.matrix.data, self.matrix.dims), (U.array.data, src_dims))
        out = U.space.range_under(self).from_numpy(out, extended_coord_data=U.extended_coord_data)
    else:
        out = U.space.range_under(self).from_xarray(unlabel_dims(self.matrix.dot(U.array_labeled_as_source())))
    if self.name != self.__class__.__name__: out.name = self.name
    return out

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def apply_adjoint(
    self:XarrayMatrixOperator, 
    V, # `XarrayVectorArray` of vectors to which the operator is applied
    mu=None # The parameter values for which to evaluate the operator
)->XarrayVectorArray: # `XarrayVectorArray` in the source `XarrayVectorSpace`
    """Apply the adjoint of the operator to an `XarrayVectorArray` in the range vector space."""
    return self.H.apply(V, mu=mu)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def _assemble_lincomb(self:XarrayMatrixOperator, operators, coefficients, identity_shift=0., name=None):
    if not all(isinstance(op, XarrayMatrixOperator) for op in operators):
        return None

    common_mat_dtype = reduce(np.promote_types,
                              (op.matrix.dtype for op in operators if hasattr(op, 'matrix')))
    common_coef_dtype = reduce(np.promote_types, (type(c) for c in coefficients + [identity_shift]))
    common_dtype = np.promote_types(common_mat_dtype, common_coef_dtype)

    if coefficients[0] == 1:
        matrix = operators[0].matrix.astype(common_dtype)
    else:
        matrix = operators[0].matrix * coefficients[0]
        if matrix.dtype != common_dtype:
            matrix = matrix.astype(common_dtype)

    for op, c in zip(operators[1:], coefficients[1:], strict=True):
        if c == 1:
            try:
                matrix += op.matrix
            except NotImplementedError:
                matrix = matrix + op.matrix
        elif c == -1:
            try:
                matrix -= op.matrix
            except NotImplementedError:
                matrix = matrix - op.matrix
        else:
            try:
                matrix += (op.matrix * c)
            except NotImplementedError:
                matrix = matrix + (op.matrix * c)

    if identity_shift: raise NotImplementedError

    return XarrayMatrixOperator(matrix, source=self.source, range=self.range)

# %% ../../nbs/api/pymor/operators.ipynb
@match_class(XarrayMatrixOperator)
def to_matrix_XarrayMatrixOperator(self, op):
    """Return the operator as a 2D matrix with stacked range and source dimensions in either `ndarray` or `sparray` format."""
    data = op.matrix.stack(_range=op.labeled_range.coords.dims, _source=op.labeled_source.coords.dims).data
    if isinstance(data, SparseArray): return sparse2d(data)
    return data
    
ToMatrixRules.insert_rule(-1, to_matrix_XarrayMatrixOperator)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def to_numpy(self:XarrayMatrixOperator):
    """Return the operator in dense `ndarray` format."""
    return self.matrix.to_numpy()

# %% ../../nbs/api/pymor/operators.ipynb
class DensifyRules(RuleTable):
    """|RuleTable| for the :func:`densify` algorithm."""

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

    @match_class(XarrayMatrixOperator)
    def action_XarrayMatrixOperator(self, op):
        op = self.replace_children(op)
        if op.sparse:
            matrix = op.matrix.copy()
            matrix.data = matrix.data.todense()
            op = op.with_(matrix=matrix)
        return op
    
    @match_class(Model, Operator)
    def action_recurse(self, op):
        return self.replace_children(op)

# %% ../../nbs/api/pymor/operators.ipynb
def densify(
    obj:Model|Operator # Object to densify
)->Model|Operator: # Densified object
    """Convert sparse operators to dense."""
    return DensifyRules().apply(obj)

# %% ../../nbs/api/pymor/operators.ipynb
class XarrayFunctionalOperator(XarrayMatrixBasedOperator):
    """An `Operator` described by a `ParameterFunctional` that assembles to a `XarrayMatrixOperator`."""
    
    def __init__(self, functional:ParameterFunctional, range, source):
        self.__auto_init(locals())

    def assemble(self, mu=None): 
        return XarrayMatrixOperator(self.functional.evaluate(mu), range=self.range.dims, source=self.source.dims)

    def apply(self, U:VectorArray, mu=None):
        return self.assemble(mu).apply(U)

    @property
    def size(self):
        return self.matrix.size

# %% ../../nbs/api/pymor/operators.ipynb
class SumOperator(Operator):
    """An `Operator` that sums over one or more dimensions of a `XarrayVectorSpace`."""
    
    linear = True
    
    def __init__(
        self, 
        source_coords:Coordinates|dict, # Coordinates to sum over 
        name=None # Operator name
    ):
        self.__auto_init(locals())
        self.source = XarrayVectorSpace(source_coords)
        self.range = XarrayVectorSpace("Sum")

    def assemble(self, mu=None): return self

    def apply(self, U, mu=None):
        U = U.space.range_under(self).from_xarray(U.array.sum(self.source.dims).expand_dims('Sum'))
        if self.name != self.__class__.__name__: U.name = self.name
        return U

    def __mul__(self, other):
        if isinstance(other, Operator):
            return ProductOperator([self, other])
        return super().__mul__(other)

    def __rmul__(self, other):
        if isinstance(other, Operator):
            return ProductOperator([other, self])
        return super().__mul__(other)

    def __str__(self):
        return f"{self.name if self.name != self.__class__.__name__ else '∑'}{_mapping_str(self)}"

# %% ../../nbs/api/pymor/operators.ipynb
class ScaleOperator(Operator):
    """A scaling operator for `XarrayVectorSpace`s."""
    
    def __init__(self, array, space=None, name=None):
        if not isinstance(array, DataArray): array = DataArray(array, coords=space.coords, name=name)
        if space is None: space = XarrayVectorSpace(array)
        self.__auto_init(locals())
        self.source = space
        self.range = space

    def assemble(self, mu=None): return self

    def apply(self, U, mu=None):
        U = U * self.array
        if self.name != self.__class__.__name__: U.name = self.name
        return U

    @property
    def display_name(self):
        return self.name if self.name != self.__class__.__name__ else 'S'
    
    def __mul__(self, other):
        if isinstance(other, Operator):
            return ProductOperator([self, other])
        return super().__mul__(other)

    def __rmul__(self, other):
        if isinstance(other, Operator):
            return ProductOperator([other, self])
        return super().__mul__(other)

    def __str__(self):
        return f"{self.display_name}{_mapping_str(self)}"

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def __init__(self:LincombOperator, operators, coefficients, solver=None, name=None):
    assert len(operators) > 0
    assert len(operators) == len(coefficients)
    assert all(isinstance(op, Operator) for op in operators)
    assert all(isinstance(c, (ParameterFunctional, Number)) for c in coefficients)
    # assert all(op.source == operators[0].source for op in operators[1:])
    # assert all(op.range == operators[0].range for op in operators[1:])
    operators = tuple(operators)
    coefficients = tuple(coefficients)

    self.operators = operators
    self.coefficients = coefficients
    self.solver = solver
    self.name = name
    try:
        self.source = max(self.operators, key=lambda o: o.source.ndim).source
        self.range = max(self.operators, key=lambda o: o.range.ndim).range
    except AttributeError:
        self.source = self.operators[0].source
        self.range = self.operators[0].range
    self.linear = all(op.linear for op in operators)


# %% ../../nbs/api/pymor/operators.ipynb
def _numpy_zero_operator(shape):
    return ZeroOperator(*[NumpyVectorSpace(dim) for dim in shape])

# %% ../../nbs/api/pymor/operators.ipynb
def _read_sparse_dataarray(
    file_name:str|Path, # File name
)->DataArray:
    if not Path(file_name).suffix: file_name = Path(file_name).with_suffix('.mtxn')
    header = []
    with open(file_name, "r") as reader:
        while True:
            line = next(reader)
            if line[0] != "%": break
            header.append(line.strip('%\n '))
    matrix = pd.read_csv(
        file_name,
        sep=r'\s+',
        header=None,
        skiprows=len(header),
        index_col=None,
        engine="c",
        comment='%'
    )
    name = header[0]
    parameters = header[1].split('\t')
    shape = [int(s) for s in header[2].split()]
    coords = {dim: index.split('\t') for dim, index in zip(header[3::2], header[4::2])}
    matrix = sparse([matrix.iloc[:, :-1].T, matrix.iloc[:, -1]], shape=shape)
    return DataArray(matrix, coords, name=name, attrs={'parameters': parameters})

# %% ../../nbs/api/pymor/operators.ipynb
@patch(cls_method=True)
@delegates(LincombOperator.__init__)
def from_file(
    cls:LincombOperator, 
    file_name:str|Path, # File name
    **kwargs # Additional arguments are passed to `LincombOperator`.
):
    """Read a `LincombOperator` (list of sparse arrays and corresponding `ExpressionParameterFunctionals`) from a .mtxn file."""
    da = _read_sparse_dataarray(file_name)
    coefficients = [
        ExpressionParameterFunctional(s, {p: 1 for p in da.parameters if p in s}) if da.parameters != [''] else 1
        for s in da.coords[da.dims[0]].data
    ]
    del da.attrs['parameters']
    # return da
    operators = [XarrayMatrixOperator(d.drop_vars(da.dims[0])) for d in da]    
    return LincombOperator(operators, coefficients, **kwargs)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def partial_evaluate(
    self:ExpressionParameterFunctional, 
    mu:dict|Mu # Parameter values 
)->float|ExpressionParameterFunctional:
    """Substitute parameter values, returning a new `ExpressionParameterFunctional` if expression does not evaluate to a number."""
    mu = {k: v[0] for k, v in mu.items()}
    expr = sy.sympify(self.expression).subs(mu)
    unevaluated = [str(s) for s in expr.free_symbols]
    if expr.is_number:
        return float(expr)
    return ExpressionParameterFunctional(
        str(expr), 
        {k: v for k, v in self.parameters.items() if k in unevaluated}
    )

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def partial_evaluate_coefficients(
    self:LincombOperator, 
    mu:Mu # Parameter values to substitute
):
    """Substitute parameter values into linear coefficients, returning a new `ExpressionParameterFunctional` if expression does not evaluate to a number."""
    result = []
    for c in self.coefficients:
        try:
            result.append(c.evaluate(mu))
        except AttributeError:
            result.append(c)
        except AssertionError:
            try:
                result.append(c.partial_evaluate(mu))
            except Exception as e:
                result.append(c)
    return result

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def partial_assemble(
    self:LincombOperator, 
    mu=None # Parameter values to substitute
)->Operator:
    """Substitute parameter values into the linear coefficients, returning a new operator with fewer (or no) parameters."""
    operators = tuple(op.assemble(mu) for op in self.operators)
    coefficients = self.partial_evaluate_coefficients(mu)
    evaluated = [[], []]
    unevaluated = [[], []]
    for o, c in zip(operators, coefficients):
        if isinstance(c, Number):
            if c:
                evaluated[0].append(o)
                evaluated[1].append(c)
        else:
            unevaluated[0].append(o)
            unevaluated[1].append(c)
    op = []
    if len(evaluated[0]):
        op.append(assemble_lincomb(
            *evaluated, 
            # solver_options=self.solver_options,
            name=self.name + '_partially_assembled'
        ))
    if len(unevaluated[0]):
        op.append(LincombOperator(*unevaluated))
    if len(op) == 0:
        op = [ZeroOperator(range=self.range, source=self.source)]
    op = np.sum(op)
    # To avoid infinite recursions, only use the result if at least one of the following
    # is true:
    #   - The operator is parametric, so the the result of assemble *must* be a different,
    #     non-parametric operator.
    #   - One of self.operators changed by calling 'assemble' on it.
    #   - The result of assemble_lincomb is of a different type than the original operator.
    #   - assemble_lincomb could simplify the list of assembled operators,
    #     which we define to be the case when the number of operators has ben reduced.
    if (self.parametric
            or operators != self.operators  # for this comparison to work self.operators always has to be a tuple!
            or type(op) != type(self)
            or len(op.operators) < len(operators)):
        return op
    else:
        return self

# %% ../../nbs/api/pymor/operators.ipynb
@patch(as_prop=True)
def terms(self:LincombOperator):
    return [c * o for c, o in zip(self.coefficients, self.operators)]

# %% ../../nbs/api/pymor/operators.ipynb
class ProductOperator(Operator):
    """An Operator given by the direct product of operators `operators`."""
    
    suboperator_symbol = '⛒'
    
    def __init__(
        self, 
        operators, # Sequence of operators that each assemble to XarrayMatrixOperator
        name=None, # Optional name
    ):
        for op in operators:
            if not isinstance(op, Operator): raise ValueError(f'Input {op} should be an Operator.')
        self.source = np.prod([op.source for op in operators])
        self.range = np.prod([op.range for op in operators])
        self.__auto_init(locals())
        self.linear = True
        
    def apply(self, U, mu=None):
        pass

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def __mul__(self:Operator, other):
    if isinstance(other, Operator) and all(isinstance(space, XarrayVectorSpace) for space in [self.source, self.range, other.source, other.range]):
        return ProductOperator([self, other])
    if isinstance(other, (Number, ParameterFunctional)):
        if self.name != 'LincombOperator' or not isinstance(self, LincombOperator):
            return LincombOperator((self,), (other,))
        else:
            return self.with_(coefficients=tuple(c * other for c in self.coefficients))
    else:
        raise TypeError(f'Object {other} is not an Operator, Number, or ParameterFunctional.')

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def __str__(self:ProductOperator):
    return ' ⨂ '.join(str(o) for o in self.operators)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def __str__(self:IdentityOperator):
    return f'I{_mapping_str(self)}'

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def apply(self:ProductOperator, U, mu=None):
    U = U.copy()
    for op in self.operators[::-1]:
        U = op.assemble(mu).apply(U)
    return U.space.range_under(self).from_xarray(U.array)

# %% ../../nbs/api/pymor/operators.ipynb
@patch(as_prop=True)
def H(self:ProductOperator):
    return self.with_(operators=[op.H for op in self.operators])

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def apply_adjoint(self:ProductOperator, V, mu=None):
    return self.H.apply(V, mu=mu)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def assemble(self:ProductOperator, mu=None):
    operators = [op.assemble(mu) for op in self.operators]
    if any(isinstance(op, ZeroOperator) for op in operators): return ZeroOperator(self.range, self.source)
    if all(op == old_op for op, old_op in zip(operators, self.operators)): return self 
    return self.with_(operators=operators)

# %% ../../nbs/api/pymor/operators.ipynb
@match_class(ProductOperator)
def expand_ProductOperator(self, op):
    op = self.replace_children(op)

    # merge child ProductOperators
    if any(isinstance(o, ProductOperator) for o in op.operators):
        ops = []
        for o in op.operators:
            if isinstance(o, ProductOperator):
                ops.extend(o.operators)
            else:
                ops.append(o)
        op = op.with_(operators=ops)

    # expand products with LincombOperators
    if any(isinstance(o, LincombOperator) for o in op.operators):
        i = next(iter(i for i, o in enumerate(op.operators) if isinstance(o, LincombOperator)))
        left, right = op.operators[:i], op.operators[i+1:]
        ops = [ProductOperator([*left, o, *right]) for o in op.operators[i].operators]
        op = op.operators[i].with_(operators=ops)

        # there can still be LincombOperators within the summands so we recurse ..
        op = self.apply(op)
    return op

#| export
ExpandRules.insert_rule(-1, expand_ProductOperator)

# %% ../../nbs/api/pymor/operators.ipynb
@match_class(IdentityOperator)
def expand_IdentityOperator(self, op):
    if not isinstance(op.source, XarrayVectorSpace): return op
    return ProductOperator([IdentityOperator(XarrayVectorSpace({k:v})) for k,v in op.space.coord_data.items()])

#| export
ExpandRules.insert_rule(-1, expand_IdentityOperator)

# %% ../../nbs/api/pymor/operators.ipynb
@match_class(ProductOperator)
def contract_ProductOperator(self, op):
    op = self.replace_children(op)
    if any(isinstance(o, ZeroOperator) for o in op.operators): return ZeroOperator(op.range, op.source)
    xarray_operators = [o for o in op.operators if     isinstance(o, XarrayMatrixOperator)]
    other_operators  = [o for o in op.operators if not isinstance(o, XarrayMatrixOperator)]
    if xarray_operators:
        matrices, ranges, sources = unzip([(o.matrix, o.range, o.source) for o in xarray_operators])
        op = XarrayMatrixOperator(
            math.prod(matrices), 
            source=math.prod(sources, start=XarrayVectorSpace()), 
            range=math.prod(ranges, start=XarrayVectorSpace())
        )
    else:
        op = 1
    if other_operators:
        op *= math.prod(other_operators) 
    return op

ContractRules.insert_rule(-2, contract_ProductOperator)

# %% ../../nbs/api/pymor/operators.ipynb
@match_class(ProductOperator)
def to_matrix_ProductOperator(self, op):
    return kron(*[self.apply(o) for o in op.operators])
    
ToMatrixRules.insert_rule(-2, to_matrix_ProductOperator)

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def as_source_array(self:ProductOperator, mu=None):
    return contract(self.assemble(mu)).as_source_array()

# %% ../../nbs/api/pymor/operators.ipynb
@patch
def as_range_array(self:ProductOperator, mu=None):
    return contract(self.assemble(mu)).as_range_array()
