"""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', 'floquet_reconstruction', 'StationaryFloquetModel']

# %% ../../nbs/api/pymor/models.ipynb
from fastcore.basics import patch, first

from pymor.basic import StationaryModel, InstationaryModel, Mu
from pymor.vectorarrays.interface import VectorArray
from ..utilities.sparse import sparse_diag, sparse_identity

# %% ../../nbs/api/pymor/models.ipynb
def _collapse_dicts(dicts):
    return {k: [d[k] for d in dicts] for k in dicts[0]}

# %% ../../nbs/api/pymor/models.ipynb
class StationaryModel(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, 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,
        )

# %% ../../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(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
from ..utilities.numpy import sym_range
from ..utilities.sparse import sparse_toeplitz
from .parameters import simplify_functionals
from .vectorarrays import XarrayVectorSpace
from .operators import XarrayMatrixOperator, XarrayFunctionalOperator
from pymor.basic import ExpressionParameterFunctional, GenericParameterFunctional
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.parsing.sympy_parser import parse_expr

# %% ../../nbs/api/pymor/models.ipynb
def _fourier_coefficients(expr:ExpressionParameterFunctional, n:int, mu=None, t:str='t', nu:str='nu'):
    """Coefficients of complex Fourier series of expr in variable `t` at frequency `nu`."""
    if isinstance(expr, ExpressionParameterFunctional): expr = expr.expression
    expr = parse_expr(expr)
    if mu is not None: expr = expr.subs(mu)
    t = sy.Symbol(t)
    series = sy.fourier_series(expr.subs(nu, 1), (t, -sy.pi, sy.pi)).truncate(n=n+1).rewrite(sy.exp).expand()
    return array([series.coeff(sy.exp(sy.I*t), n=k) if k != 0 else series.coeff(t, n=0) for k in range(-n, n+1)]).astype(complex)

# %% ../../nbs/api/pymor/models.ipynb
def _fourier_range_coord(n): return first(Coordinates({'Fourier coefficient (range)': sym_range(n)}).values())
def _fourier_source_coord(n): return first(Coordinates({'Fourier coefficient': sym_range(n)}).values())
def _fourier_operator_coords(n): return (_fourier_range_coord(n) * _fourier_source_coord(n)).coords
def _fourier_range_space(n): return XarrayVectorSpace(_fourier_range_coord(n))
def _fourier_source_space(n): return XarrayVectorSpace(_fourier_source_coord(n))

# %% ../../nbs/api/pymor/models.ipynb
def _fourier_expansion_operator(coefficient, n_operator, n_vector, mu=None, t:str='t', nu:str='nu'):
    return XarrayMatrixOperator(
        DataArray(
            sparse_toeplitz(np.pad(_fourier_coefficients(coefficient, n=n_operator, mu=mu), n_vector + 1)), 
            coords=_fourier_operator_coords(n_vector)
        )
    )

# %% ../../nbs/api/pymor/models.ipynb
def _fourier_identity_operator(n_vector):
    return XarrayMatrixOperator(
        DataArray(
            sparse_identity(2*n_vector+1), 
            coords=_fourier_operator_coords(n_vector)
        )
    )    

# %% ../../nbs/api/pymor/models.ipynb
def _fourier_diagonal_operator(vec):
    return XarrayMatrixOperator(
        DataArray(
            sparse_diag(vec),
            coords=_fourier_operator_coords((len(vec)-1)//2)
        )
    )    

# %% ../../nbs/api/pymor/models.ipynb
def _dm_identity_operator(lgs): 
    return XarrayMatrixOperator(DataArray(sparse_identity(lgs.n_variables), lgs.A_ind.operators[0].matrix.coords))

# %% ../../nbs/api/pymor/models.ipynb
def _dm_identity_operator(lgs): 
    return XarrayMatrixOperator(DataArray(sparse_identity(lgs.n_variables), lgs.A_ind.operators[0].matrix.coords))

# %% ../../nbs/api/pymor/models.ipynb
def _da_identity(range_coord, source_coord):
    return DataArray(sparse_identity(len(source_coord)), coords=((range_coord.name, range_coord.data), (source_coord.name, source_coord.data)))

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

# %% ../../nbs/api/pymor/models.ipynb
def _identity_operator(range, source):
    return math.prod(XarrayMatrixOperator(_da_identity(r, s)) for r, s in zip(range.coords.values(), source.coords.values()))

# %% ../../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_coord(time): return first(Coordinates({'Time': time}).values())
def _time_space(time): return XarrayVectorSpace(_time_coord(time))

# %% ../../nbs/api/pymor/models.ipynb
def _floquet_reconstruction_array(n_vector, time, mu):
    return np.exp(
            1j * ExpressionParameterFunctional('nu', {'nu': 1}).evaluate(mu)
            * _time_coord(time)
            * _fourier_source_coord(n_vector)
        )

# %% ../../nbs/api/pymor/models.ipynb
def floquet_reconstruction(n_vector, time):
    return XarrayFunctionalOperator(
        GenericParameterFunctional(lambda mu: _floquet_reconstruction_array(n_vector, time, mu), {'nu': 1}),
        range=_time_space(time),
        source=_fourier_source_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
        nu:'nu', # Name of the frequency variable
        n_operator=None, # Number of terms in the operator Fourier series. `None` defaults to `n_vector` - 1
        products=None,
        name=None,
        data=None
    ):
        if n_operator is None: n_operator = n_vector - 1
        operator = expand(operator)
        A = sum(term for term in operator.terms if not set(mu).intersection(term.parameters))
        B = sum(term for term in operator.terms if     set(mu).intersection(term.parameters))
        B = simplify_functionals(B)
        B_expanded = sum(_fourier_expansion_operator(c, n_operator, n_vector, mu=mu) * o for c, o in zip(B.coefficients, B.operators))
        A_expanded = _fourier_identity_operator(n_vector) * A 
        D_expanded = (
            ExpressionParameterFunctional(f'-(1j)*{nu}', {nu: 1})
            * (_fourier_diagonal_operator(sym_range(n_vector)) 
               * _identity_operator(operator.range, operator.source)
              )
        )
        operator_expanded = (D_expanded + A_expanded + B_expanded).with_(solver_options=operator.solver_options)
        rhs_expanded = _kron_vector(_fourier_range_space(n_vector)) * rhs
        self.mu, self.n_vector, self.nu, self.n_operator = mu, n_vector, nu, n_operator
        super().__init__(
            operator_expanded,
            rhs_expanded,
            products=products,
            data=data
        )
