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

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

# %% auto 0
__all__ = ['init_progress_widget', 'tick', 'update_progress_widget', 'init_cvode', 'solve_cvode', 'AdamsTimeStepper',
           'BDFTimeStepper', 'StiffSwitchingTimeStepper']

# %% ../../nbs/api/pymor/timestepping.ipynb
import time

from fastcore.basics import patch
import ipywidgets as widgets

import numpy as np
from scikits.odes.sundials import cvode
import scipy.sparse as sps
import pymor
from pymor.algorithms.to_matrix import to_matrix
from pymor.algorithms.simplify import expand
from pymor.algorithms.timestepping import TimeStepper
from pymor.basic import LincombOperator

from pymor.algorithms.simplify import ExpandRules, expand, ContractRules, contract
from pymor.algorithms.rules import match_class, RuleTable


from ..utilities.sparse import sparse2d_identity, spilu, restrict_bandwidth
from .vectorarrays import *
from .operators import *
# from pymor.basic 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=100000,
        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
def init_progress_widget(initial_time, end_time):
    progress = widgets.FloatProgress(
        value=initial_time,
        min=initial_time,
        max=end_time,
        bar_style='info',
        orientation='horizontal'
    )
    step_size = widgets.Label()
    widget = widgets.HBox([progress, step_size])
    display(widget)
    return widget

# %% ../../nbs/api/pymor/timestepping.ipynb
def tick(): return time.time() * 1000

# %% ../../nbs/api/pymor/timestepping.ipynb
def update_progress_widget(widget, tock, t, last_t):
    newtick = tick()
    if newtick - tock < 10: return tock
    widget.children[0].value = t
    widget.children[1].value = f'{t - last_t:1.1e}'
    return newtick

# %% ../../nbs/api/pymor/timestepping.ipynb
def init_cvode(cvode_rhs, num_values):
    return cvode.CVODE(cvode_rhs, lmm_type='Adams', nonlinsolver='fixedpoint', max_steps=1000000, one_step_compute=num_values is None)

# %% ../../nbs/api/pymor/timestepping.ipynb
def solve_cvode(solver, t_list, initial_data):
    return solver.solve(t_list, initial_data.to_numpy().T[0])

# %% ../../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):

        def cvode_rhs(t, y, ydot):
            y = operator.source.from_numpy(y)
            np.copyto(ydot, rhs.as_range_array().to_numpy().ravel() - operator.apply(y, mu.at_time(t=t)).to_numpy().ravel())

        self._solver = cvode.CVODE(
            cvode_rhs, 
            lmm_type='Adams', 
            nonlinsolver='fixedpoint', 
            max_steps=1000000, 
            one_step_compute=True
        )
        min_reportable_step = (end_time - initial_time)/num_values
        progress = init_progress_widget(initial_time, end_time)
        tock = tick()
        y = [self._solver.init_step(initial_time, initial_data.to_numpy().T[0]).values.y]
        t = [initial_time]
        last_t = initial_time
        while t[-1] < end_time:
            values = self._solver.step(end_time).values
            tock = update_progress_widget(progress, tock, values.t, last_t)
            last_t = values.t
            if values.t - t[-1] < min_reportable_step: continue
            t.append(values.t)
            y.append(values.y)
        t = np.array(t)
        y = np.array(y)
        progress.close()
        
        if isinstance(operator.source, XarrayVectorSpace):
            return [operator.source.from_numpy(y, extended_coord_data={'Time': t})]
        return ((operator.source.from_numpy(u), t) for u, t in zip(y, t))
            
        # self._solver = init_cvode(cvode_rhs, num_values)
        # t_list = np.linspace(initial_time, end_time, num_values)
        # sol = solve_cvode(self._solver, t_list, initial_data)
        # progress.close()
        # if isinstance(operator.source, XarrayVectorSpace):
        #     return [operator.source.from_numpy(sol.values.y, extended_coord_data={'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 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):
        options = cvode_solver_options()['cvode_bdf']
        # if solver_options:
        #     options.update(solver_options)
        
        def cvode_rhs(t, y, ydot):
            np.copyto(ydot, (-operator.assemble(mu.at_time(t=t)).apply(operator.source.from_numpy(y)) + rhs.as_range_array(mu.at_time(t=t))).to_numpy().T[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.at_time(t=t)))
                # if operator.solver_options is not None and 'inverse' in operator.solver_options:
                #     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])
            if sps.issparse(p): 
                p = p.tocsc()
            else:
                p = sps.csc_array(p)
            user_data['factored_preconditioner'] = spilu(p) # , 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={},
            one_step_compute=True
        )
        # t_list = np.linspace(initial_time, end_time, num_values)
        progress = init_progress_widget(initial_time, end_time)
        tock = tick()
        # last_t = 0
        y = [self._solver.init_step(initial_time, initial_data.to_numpy().T[0]).values.y]
        t = [initial_time]
        while t[-1] < end_time:
            values = self._solver.step(end_time).values
            t.append(values.t)
            y.append(values.y)
            tock = update_progress_widget(progress, tock, t[-1], t[-2])
        t = np.array(t)
        y = np.array(y)
        # sol = self._solver.solve(t_list, initial_data.to_numpy().T[0])
        progress.close()
        
        if isinstance(operator.source, XarrayVectorSpace):
            return [operator.source.from_numpy(y, extended_coord_data={'Time': t})]
        return ((operator.source.from_numpy(u), t) for u, t in zip(y, t))

# %% ../../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
