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

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

# %% auto 0
__all__ = ['AdamsTimeStepper', 'BDFTimeStepper', 'StiffSwitchingTimeStepper']

# %% ../../nbs/api/pymor/timestepping.ipynb
from fastcore.basics import patch
import ipywidgets as widgets

import numpy as np
from scikits.odes.sundials import cvode
import pymor
from pymor.algorithms.to_matrix import to_matrix
from pymor.algorithms.timestepping import TimeStepper

from ..utilities.sparse import sparse2d_identity, spilu, restrict_bandwidth
from .vectorarrays import *
from .operators import *

# %% ../../nbs/api/pymor/timestepping.ipynb
try:
    @pymor.defaults('cvode_bdf_rtol', 'cvode_bdf_atol', 'cvode_bdf_max_steps', 'cvode_stiff_switching_t_switch')
    def cvode_solver_options(
        cvode_bdf_rtol=1e-3,
        cvode_bdf_atol=1e-6,
        cvode_bdf_max_steps=1000,
        cvode_bdf_inflection_times=None,
        cvode_stiff_switching_t_switch=1e-7
    ):
        opts = {
            'cvode_bdf': {'type': 'cvode_bdf', 'atol': cvode_bdf_atol, 'rtol': cvode_bdf_rtol, 'max_steps': cvode_bdf_max_steps, 'inflection_times': cvode_bdf_inflection_times},
            'cvode_stiff_switching': {'type': 'cvode_stiff_switching', 't_switch': cvode_stiff_switching_t_switch}
        }
        return opts
except ValueError: pass

# %% ../../nbs/api/pymor/timestepping.ipynb
@patch
def solve(self:TimeStepper, initial_time, end_time, initial_data, operator, rhs=None, mass=None, mu=None, num_values=None):
    """Apply time-stepper to the equation.

    The equation is of the form ::

        M(mu) * d_t u + A(u, mu, t) = F(mu, t),
                         u(mu, t_0) = u_0(mu).

    Parameters
    ----------
    initial_time
        The time at which to begin time-stepping.
    end_time
        The time until which to perform time-stepping.
    initial_data
        The solution vector at initial_timeinitial_time.
    operator
        The |Operator| A.
    rhs
        The right-hand side F (either |VectorArray| of length 1 or |Operator| with
        source.dim == 1source.dim == 1). If NoneNone, zero right-hand side is assumed.
    mass
        The |Operator| M. If NoneNone, the identity operator is assumed.
    mu
        |Parameter values| for which operatoroperator and rhsrhs are evaluated. The current
        time is added to mumu with key tt.
    num_values
        The number of returned vectors of the solution trajectory. If NoneNone, each
        intermediate vector that is calculated is returned.

    Returns
    -------
    |VectorArray| containing the solution trajectory.
    """
    try:
        num_time_steps = self.estimate_time_step_count(initial_time, end_time)
    except NotImplementedError:
        num_time_steps = 0
    iterator = self.iterate(initial_time, end_time, initial_data, operator, rhs=rhs, mass=mass, mu=mu,
                            num_values=num_values)
    iterator = list(iterator)
    if isinstance(iterator[0], XarrayVectorArray):
        return iterator[0]
    U = operator.source.empty(reserve=num_values if num_values else num_time_steps + 1)
    t = []
    for U_n, t_n in iterator:
        U.append(U_n)
        t.append(t_n)
    try: return operator.source.from_numpy(U.to_numpy(), l={"Time": t})
    except: return U

# %% ../../nbs/api/pymor/timestepping.ipynb
class AdamsTimeStepper(TimeStepper):
    def __init__(self):
        self.__auto_init(locals())
        
    def iterate(self, initial_time, end_time, initial_data, operator, rhs=None, mass=None, mu=None, num_values=None):
        a = operator.assemble(mu.with_(t=0.)).matrix
        b = rhs.to_numpy().ravel()
        
        def cvode_rhs(t, y, ydot):
            np.copyto(ydot, (-a.dot(y) + b))
        
        self._solver = cvode.CVODE(cvode_rhs, lmm_type='Adams', nonlinsolver='fixedpoint', max_steps=1000000, one_step_compute=num_values is None)
        if num_values is not None:
            self._t_list = np.linspace(initial_time, end_time, num_values)
            sol = self._solver.solve(self._t_list, initial_data.to_numpy()[0])
            y = sol.values.y
        else:
            self._t_list = [initial_time]
            y = list(initial_data.to_numpy())
            self._solver.init_step(t0=self._t_list[0], y0=y[0])
            while self._t_list[-1] < end_time:
                sol = self._solver.step(t=end_time)
                self._t_list.append(sol.values.t)
                y.append(sol.values.y)
        return ((operator.range.from_numpy(u), t) for u, t in zip(y, self._t_list))

# %% ../../nbs/api/pymor/timestepping.ipynb
class BDFTimeStepper(TimeStepper):
    def __init__(self):
        # self.__auto_init(locals())
        pass
        
    def iterate(self, initial_time, end_time, initial_data, operator, rhs=None, mass=None, mu=None, num_values=None, solver_options=None):
        options = cvode_solver_options()['cvode_bdf']
        if solver_options:
            options.update(solver_options)

        progress = widgets.FloatProgress(
            value=initial_time,
            min=initial_time,
            max=end_time,
            bar_style='info',
            orientation='horizontal'
        )
        display(progress)
        
        def cvode_rhs(t, y, ydot):
            progress.value = t
            np.copyto(ydot, (-operator.assemble(mu.with_(t=t)).apply(operator.source.from_numpy(y)) + rhs.as_range_array(mu.with_(t=t))).to_numpy()[0])

        def preconditioner_setup(t, y, jok, jcurPtr, gamma, user_data):
            """Generate P and do ILU decomposition."""
            if jok:
                jcurPtr.value = False
            else:
                user_data['approximate_jacobian'] = -to_matrix(operator.assemble(mu.with_(t=t)))
                user_data['approximate_jacobian'] = restrict_bandwidth(user_data['approximate_jacobian'], operator.solver_options['inverse']['preconditioner_bandwidth'])
                jcurPtr.value = True
            # Scale jacobian by -gamma, add identity matrix and do LU decomposition
            p = -gamma*user_data['approximate_jacobian'] + sparse2d_identity(user_data['approximate_jacobian'].shape[0])
            user_data['factored_preconditioner'] = spilu(p.tocsc()) # , permc_spec='NATURAL')
            return 0

        def preconditioner_solve(t, y, r, z, gamma, delta, lr, user_data):
            """ Solve the block-diagonal system Pz = r. """
            np.copyto(z, user_data['factored_preconditioner'].solve(r))
            return 0              
        
        self._solver = cvode.CVODE(
            cvode_rhs,
            lmm_type='BDF', 
            nonlinsolver='newton', 
            linsolver='spgmr',
            precond_type='left',
            prec_setupfn=preconditioner_setup, 
            prec_solvefn=preconditioner_solve,
            rtol=options['rtol'], 
            atol=options['atol'], 
            max_steps=options['max_steps'],
            user_data={}
        )
        t_list = np.linspace(initial_time, end_time, num_values)
        sol = self._solver.solve(t_list, initial_data.to_numpy()[0])
        progress.close()
        if isinstance(operator.source, XarrayVectorSpace):
            return [operator.source.from_numpy(sol.values.y, extended_dim={'Time': t_list})]
        return ((operator.source.from_numpy(u), t) for u, t in zip(sol.values.y, t_list))

# %% ../../nbs/api/pymor/timestepping.ipynb
class StiffSwitchingTimeStepper(TimeStepper):
    def __init__(self):
        # self.__auto_init(locals())
        pass
        
    def iterate(self, initial_time, end_time, initial_data, operator, rhs=None, mass=None, mu=None, num_values=None, solver_options=None):
        options = cvode_solver_options()['cvode_stiff_switching']
        if solver_options:
            options.update(solver_options)

        if initial_time < options['t_switch']:
            nonstiff_solver = AdamsTimeStepper()
            nonstiff = list(nonstiff_solver.iterate(initial_time, options['t_switch'], initial_data, operator, rhs=rhs, mass=mass, mu=mu, num_values=num_values))
            nonstiff_t_list = nonstiff_solver._t_list
            initial_data = nonstiff[-1][0]
            initial_time = nonstiff_t_list[-1]
        else:
            nonstiff = []
            nonstiff_t_list = []

        if end_time > options['t_switch']:
            stiff_solver = BDFTimeStepper()
            stiff = list(stiff_solver.iterate(options['t_switch'], end_time, initial_data, operator, rhs=rhs, mass=mass, mu=mu, num_values=num_values, solver_options=solver_options))
            stiff_t_list = stiff_solver._t_list            
        else:
            stiff = []
            stiff_t_list = []

        self._t_list = np.concatenate([nonstiff_t_list, stiff_t_list])
        
        return nonstiff + stiff
