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

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

# %% auto 0
__all__ = ['StationaryModel', 'InstationaryModel', 'fourier_space', 'to_sympy', 'fourier_series', 'fourier_series_coefficients',
           'fourier_expansion_operator', 'FourierExpandRules', 'fourier_expand', 'fourier_diagonal_operator', 'd_term',
           'floquet_expand', 'time_space', 'floquet_reconstruction_array', 'floquet_reconstruction_operator',
           'StationaryFloquetModel']

# %% ../../nbs/api/pymor/models.ipynb
from functools import partial
from numbers import Number

from fastcore.basics import patch, first
import pymor
from pymor.basic import Mu, LincombOperator, IdentityOperator
from pymor.vectorarrays.interface import VectorArray
from ..utilities.sparse import sparse_diag, sparse_identity

# %% ../../nbs/api/pymor/models.ipynb
class StationaryModel(pymor.basic.StationaryModel):
    """Extend pyMOR's `StationaryModel` to include a `data` attribute for storing, e.g., the model's grid."""
    def __init__(
        self,
        operator,
        rhs,
        output_functional=None,
        products=None,
        error_estimator=None,
        visualizer=None,
        name=None,
        data=None
    ):
        self.data = data
        if products:
            for k, v in products.items():
                setattr(self, k, partial(self._apply_product, v=v))
                # setattr(self, k, lambda u, mu=None, v=v: v.apply(u, Mu(mu) if mu is not None else None) if isinstance(u, VectorArray) else v.apply(self.solve(u), mu=Mu(u)))
        super().__init__(
            operator,
            rhs,
            output_functional=output_functional,
            products=products,
            error_estimator=error_estimator,
            visualizer=visualizer,
            name=name,
        )

    def _apply_product(self, u, mu=None, v=None):
        if mu is not None: mu = Mu(mu)
        if isinstance(u, VectorArray):
            return v.apply(u, mu)
        return v.apply(self.solve(u), mu=mu)

# %% ../../nbs/api/pymor/models.ipynb
@patch
def solve(self:StationaryModel, mu=None, input=None, return_error_estimate=False, **kwargs):
    """Extend `StationaryModel.solve` to solve over a range of parameters."""
    if not isinstance(mu, (list, tuple)):
        data = self.compute(
            solution=True,
            solution_error_estimate=return_error_estimate,
            mu=mu,
            input=input,
            **kwargs
        )
        if return_error_estimate:
            return data['solution'], data['solution_error_estimate']
        else:
            return data['solution']
    else:
        fixed_mu, var_mu = mu
        n = len(first(var_mu.values()))
        if not all(len(v) == n for v in var_mu.values()): raise ValueError(f'All varying parameter lists {var_mu} must be the same length.')
        sol = [self.solve(fixed_mu | {k: v[i] for k, v in var_mu.items()}, input=input, return_error_estimate=False, **kwargs) for i in range(n)]
        return self.operator.source.concatenate(sol, var_mu)

# %% ../../nbs/api/pymor/models.ipynb
class InstationaryModel(pymor.basic.InstationaryModel):
    """Extend pyMOR's `InstationaryModel` to include a `data` attribute for storing, e.g., the model's grid."""
    def __init__(
        self,
        T,
        initial_data,
        operator,
        rhs,
        mass=None,
        time_stepper=None,
        num_values=None,
        output_functional=None,
        products=None,
        error_estimator=None,
        visualizer=None,
        name=None,
        data=None
    ):
        self.data = data
        if products:
            for k, v in products.items():
                setattr(self, k, lambda u, mu=None, v=v: v.apply(u, mu))
        super().__init__(
        T,
        initial_data,
        operator,
        rhs,
        mass=mass,
        time_stepper=time_stepper,
        num_values=num_values,
        output_functional=output_functional,
        products=products,
        error_estimator=error_estimator,
        visualizer=visualizer,
        name=name,
    )

# %% ../../nbs/api/pymor/models.ipynb
import math

from ..utilities.numpy import sym_range
from ..utilities.sparse import sparse_toeplitz, dense_toeplitz
from .parameters import simplify_functionals
from .vectorarrays import XarrayVectorSpace
from .operators import XarrayMatrixOperator, XarrayFunctionalOperator
from pymor.basic import ExpressionParameterFunctional, GenericParameterFunctional
from pymor.parameters.functionals import ProductParameterFunctional
from pymor.algorithms.simplify import expand
from xarray import Coordinates, DataArray
import numpy as np
from numpy import array
import sympy as sy
from sympy import Expr
from sympy.parsing.sympy_parser import parse_expr

# %% ../../nbs/api/pymor/models.ipynb
def fourier_space(n): 
    return XarrayVectorSpace({'Fourier coefficient': sym_range(n)})

# %% ../../nbs/api/pymor/models.ipynb
def to_sympy(functional):
    if isinstance(functional, Number):
        return functional
    if isinstance(functional, ProductParameterFunctional):
        return math.prod(to_sympy(factor) for factor in functional.factors)
    if isinstance(functional, ExpressionParameterFunctional): 
        out = functional.expression
    else:
        out = functional
    return parse_expr(out)

# %% ../../nbs/api/pymor/models.ipynb
def fourier_series(expr, n, t:str='t', nu:str='nu'):
    t = sy.Symbol(t)
    return sy.fourier_series(
        expr.subs(nu, 1), 
        (t, -sy.pi, sy.pi)
    ).truncate(
        n=n+1
    ).rewrite(
        sy.exp
    ).expand()

# %% ../../nbs/api/pymor/models.ipynb
def fourier_series_coefficients(series:Expr, n:int, t:str='t'):
    """Returns the positive and negative Fourier series coefficients up to and including order `n`.
    """
    t = sy.Symbol(t)
    return array([
        series.coeff(sy.exp(sy.I*t) if k != 0 else t, n=k) 
        for k in range(-n, n+1)
    ])

# %% ../../nbs/api/pymor/models.ipynb
def fourier_expansion_operator(coefficient, n_operator, n_vector, mu:dict=None, t:str='t', nu:str='nu') -> np.ndarray:
    """Coefficients of complex Fourier series of expr in variable `t` at frequency `nu`.

    Assumes that either the coefficients are numerical, or there is only the zero-order coefficient. Could be useful to generalize.
    """
    space = fourier_space(n_vector)
    if isinstance(coefficient, Number):
        return IdentityOperator(space) * coefficient        
    expr = to_sympy(coefficient)
    if mu is not None: expr = expr.subs(mu)
    if not expr.has(sy.Symbol(t)): 
        return IdentityOperator(space) * coefficient
    series = fourier_series(expr, n_operator, t=t, nu=nu)
    coeffs = fourier_series_coefficients(series, n_operator, t=t).astype(complex)
    return XarrayMatrixOperator(
        sparse_toeplitz(np.pad(coeffs, n_vector + 1)),
        source=space,
        range=space
    )

# %% ../../nbs/api/pymor/models.ipynb
from pymor.algorithms.rules import match_class, RuleTable
from pymor.models.interface import Model
from pymor.operators.interface import Operator

# %% ../../nbs/api/pymor/models.ipynb
class FourierExpandRules(RuleTable):
    """|RuleTable| to Fourier expand an operator."""

    def __init__(self, n_operator:int, n_vector:int, mu:dict|None=None, t:str='t', nu:str='nu'):
        super().__init__(use_caching=True)
        self.__auto_init(locals())

    @match_class(LincombOperator)
    def action_LincombOperator(self, op):
        if not set(self.mu).intersection(op.parameters): return op * IdentityOperator(fourier_space(self.n_vector))
        # op = self.replace_children(op)
        return sum(
            fourier_expansion_operator(c, self.n_operator, self.n_vector, mu=self.mu, t=self.t, nu=self.nu) * o
            for c, o in zip(op.coefficients, op.operators)
        )

    @match_class(Model, Operator)
    def action_recurse(self, op):
        if not set(self.mu).intersection(op.parameters): return op * IdentityOperator(fourier_space(self.n_vector))
        return self.replace_children(op)

# %% ../../nbs/api/pymor/models.ipynb
def fourier_expand(obj, n_operator:int, n_vector:int, mu:dict|None=None, t:str='t', nu:str='nu'):
    return FourierExpandRules(n_operator, n_vector, mu=mu).apply(obj)

# %% ../../nbs/api/pymor/models.ipynb
def fourier_diagonal_operator(vec):
    space = fourier_space((len(vec)-1)//2)
    return XarrayMatrixOperator(sparse_diag(vec), source=space, range=space)

# %% ../../nbs/api/pymor/models.ipynb
def d_term(n_vector:int, nu:str='nu'):
    return ExpressionParameterFunctional(f'-(1j)*{nu}', {nu: 1}) * fourier_diagonal_operator(sym_range(n_vector))

# %% ../../nbs/api/pymor/models.ipynb
def floquet_expand(op, n_operator, n_vector, mu:dict=None, t:str='t', nu:str='nu'):
    return (
        d_term(n_vector, nu=nu) * IdentityOperator(op.source)
        + fourier_expand(op, n_operator, n_vector, mu=mu, t=t, nu=nu)
    ).with_(solver=op.solver)

# %% ../../nbs/api/pymor/models.ipynb
def _kron_vector(space): return space.from_numpy((sym_range((space.dim - 1)//2) == 0).astype(int))

# %% ../../nbs/api/pymor/models.ipynb
def time_space(time): return XarrayVectorSpace({'Time': time})

# %% ../../nbs/api/pymor/models.ipynb
def floquet_reconstruction_array(n_vector, time, mu):
    time_coord = time_space(time).coords['Time']
    fourier_coord = fourier_space(n_vector).coords['Fourier coefficient']
    nu = ExpressionParameterFunctional('nu', {'nu': 1}).evaluate(mu)
    return np.exp(-1j * nu * time_coord * fourier_coord)

# %% ../../nbs/api/pymor/models.ipynb
def floquet_reconstruction_operator(n_vector, time):
    return XarrayFunctionalOperator(
        GenericParameterFunctional(lambda mu: floquet_reconstruction_array(n_vector, time, mu), {'nu': 1}),
        range=time_space(time),
        source=fourier_space(n_vector),
    )

# %% ../../nbs/api/pymor/models.ipynb
class StationaryFloquetModel(StationaryModel):
    """A model that solves for the Fourier coefficients of the periodic state of a modulated system."""

    def __init__(
        self,
        operator, # System evolution operator
        rhs, # System right-hand side vector
        mu, # Parameter specification for parameters modulated at a frequency `nu` 
        n_vector, # Number of terms in the system Fourier series
        time, # Time values for displaying result
        nu='nu', # Name of the frequency variable
        t='t', # Name of the time variable
        n_operator=None, # Number of terms in the operator Fourier series. `None` defaults to `n_vector` - 1
        products=None,
        name=None,
        data=None,
        solver=None,
        preconditioner=None,
    ):
        if n_operator is None: n_operator = n_vector - 1
        operator = floquet_expand(expand(operator), n_operator, n_vector, mu=mu, t=t, nu=nu)
        if solver is not None:
            if preconditioner is not None:
                preconditioner = floquet_expand(expand(preconditioner), n_operator, n_vector, mu=mu, t=t, nu=nu)
                solver = solver.with_(preconditioner=preconditioner)
            operator = operator.with_(solver=solver)
        rhs = _kron_vector(fourier_space(n_vector)) * rhs
        super().__init__(
            operator=operator,
            rhs=rhs,
            products=(products or {}) | self._floquet_products(n_vector, time),
            data=data
        )
        self.__auto_init(locals())

    @staticmethod
    def _floquet_products(n_vector, time): 
        return {
            'reconstruct': floquet_reconstruction_operator(n_vector, time),
        }
