import numpy as np
import math
from abc import ABC, abstractmethod
import scipy.sparse as sp
from .nonlinear_solvers import ImplicitEquationSolver  # Relative import for a solver class


class IntegrationMethod(ABC):
    """
    Abstract base class for integration methods.

    Classes derived from IntegrationMethod must implement the `step` method,
    which advances the solution of an ODE from time t to t+h.
    """

    @abstractmethod
    def step(self, fun, t, y, h):
        """
        Advance the solution of an ODE by one time step.

        Parameters:
            fun: callable
                The function defining the ODE (dy/dt = fun(t, y)).
            t: float
                The current time.
            y: np.array
                The current state vector.
            h: float
                The time step size.

        Returns:
            The new state after taking the step (and possibly additional diagnostic info).
        """
        pass


class BackwardEuler(IntegrationMethod):
    """
    Implements the Backward Euler implicit integration method.

    Attributes:
        solver: ImplicitEquationSolver
            A solver instance used to solve the implicit equations.
        A: np.array or None
            A matrix used in the formulation of the method. If None, the identity matrix is used.
        use_identity: bool
            Flag to indicate whether to use the identity matrix.
        _ID_CACHE: dict
            Class-level cache for identity matrices to avoid recomputation.
    """

    # Cache for identity matrices to avoid repeated allocation.
    # Keys: ('dense'| 'csr', n)
    _ID_CACHE = {}

    def __init__(self, solver=None, A=None, pass_prev_state=False, pass_step_size=False):
        """
        Initialize a Backward Euler integration method.

        Parameters:
            solver: ImplicitEquationSolver, optional
                The solver to use for solving the implicit equation. Defaults to using
                an ImplicitEquationSolver with method 'semismooth_newton'.
            A: np.array, optional
                The matrix used in the formulation. If not provided, identity matrix is used.
            pass_prev_state: bool, optional
                When True, the previously accepted state ``y`` will be supplied (when
                supported by the callable) as an additional argument to both the RHS
                function and its Jacobian.
            pass_step_size: bool, optional
                When True, the step size ``h`` for the current implicit solve will be
                forwarded to the RHS and Jacobian callables (when their signatures
                accept it).
        """
        self.solver = solver or ImplicitEquationSolver(method='semismooth_newton')
        self.A = A
        self.use_identity = (A is None)
        self.pass_prev_state = pass_prev_state
        self.pass_step_size = pass_step_size
        # Method order (for adaptive controllers)
        self.order = 1
        # Per-instance caches for bound call wrappers to avoid repeated try/except dispatch
        # Keys are tuples of (id(func), has_prev, has_h)
        self._fun_bindings = {}
        self._jac_bindings = {}

    def _get_bound_wrapper(self, func, has_prev, has_h, cache):
        """Bind a lightweight caller for func based on available arguments and cache it.

        Returns a callable wrapper(tt, yy, fk_val, prev_state, step_size) -> result
        that uses a fixed argument order discovered once.
        """
        key = (id(func), bool(has_prev), bool(has_h))
        wrapper = cache.get(key)
        if wrapper is not None:
            return wrapper

        # Candidate orders filtered by availability
        # Prefer t,y first; fall back to y-only variants
        orders = [
            ('t', 'y', 'prev', 'Fk', 'h'),
            ('t', 'y', 'prev', 'Fk'),
            ('t', 'y', 'prev', 'h'),
            ('t', 'y', 'Fk', 'h'),
            ('t', 'y', 'prev'),
            ('t', 'y', 'Fk'),
            ('t', 'y', 'h'),
            ('t', 'y'),
            ('y', 'prev', 'Fk', 'h'),
            ('y', 'prev', 'Fk'),
            ('y', 'prev', 'h'),
            ('y', 'Fk', 'h'),
            ('y', 'prev'),
            ('y', 'Fk'),
            ('y', 'h'),
            ('y',),
        ]

        def _build(order):
            def _call(tt, yy, fk, prev, h):
                # Skip labels not available
                args = []
                for lab in order:
                    if lab == 't':
                        if tt is None:
                            return _sentinel
                        args.append(tt)
                    elif lab == 'y':
                        if yy is None:
                            return _sentinel
                        args.append(yy)
                    elif lab == 'prev':
                        if not has_prev or prev is None:
                            return _sentinel
                        args.append(prev)
                    elif lab == 'Fk':
                        if fk is None:
                            return _sentinel
                        args.append(fk)
                    elif lab == 'h':
                        if not has_h or h is None:
                            return _sentinel
                        args.append(h)
                return func(*args)
            return _call

        _sentinel = object()
        # Probe using a real call at first use; we rely on the caller to give real values
        def resolve_and_cache(tt, yy, fk, prev, h):
            for order in orders:
                # quick availability filter
                if ('prev' in order and (not has_prev or prev is None)):
                    continue
                if ('h' in order and (not has_h or h is None)):
                    continue
                caller = _build(order)
                try:
                    out = caller(tt, yy, fk, prev, h)
                    if out is _sentinel:
                        continue
                    # Freeze wrapper with fixed order
                    def bound(tt2, yy2, fk2, prev2, h2, _order=order, _func=func):
                        args = []
                        for lab in _order:
                            if lab == 't':
                                args.append(tt2)
                            elif lab == 'y':
                                args.append(yy2)
                            elif lab == 'prev':
                                args.append(prev2)
                            elif lab == 'Fk':
                                args.append(fk2)
                            elif lab == 'h':
                                args.append(h2)
                        return _func(*args)
                    cache[key] = bound
                    return out
                except TypeError:
                    continue
            raise TypeError(
                "Unable to call {} with supported signatures.".format(getattr(func, '__name__', 'callable'))
            )

        # Return a wrapper that resolves on first call, then reuses cached bound
        def wrapper_first(tt, yy, fk, prev, h):
            # Attempt resolve; this stores the bound wrapper in cache
            out = resolve_and_cache(tt, yy, fk, prev, h)
            # Use cached for subsequent calls
            return out

        return wrapper_first

    def _get_A(self, n):
        """
        Retrieve or compute the appropriate matrix A of size (n x n).

        If no specific matrix A is provided during initialization, return an identity
        matrix. For large systems (as indicated by the solver's sparse preference),
        a CSR identity is returned to avoid allocating a dense n×n array.

        Returns a numpy.ndarray or scipy.sparse.csr_matrix matching the chosen path.
        """
        if not self.use_identity:
            return self.A

        # Prefer sparse identity for large n (same heuristic as solver sparse path if available)
        want_sparse = False
        try:
            want_sparse = bool(self.solver._sparse_active(n))  # may not exist for custom solvers
        except Exception:
            # Fallback heuristic using solver's threshold when present
            thr = getattr(self.solver, 'sparse_threshold', 200)
            try:
                want_sparse = (n >= int(thr))
            except Exception:
                want_sparse = (n >= 200)

        if want_sparse:
            key = ('csr', n)
            if key not in self._ID_CACHE:
                self._ID_CACHE[key] = sp.eye(n, format='csr')
            return self._ID_CACHE[key]
        else:
            key = ('dense', n)
            if key not in self._ID_CACHE:
                self._ID_CACHE[key] = np.eye(n)
            return self._ID_CACHE[key]

    def step(self, fun, t, y, h):
        """Perform one implicit Backward Euler step.

        Solves the nonlinear system::

            A ((y_new - y) / h) - f(t + h, y_new) = 0

        Parameters
        ----------
        fun : callable
            RHS function ``fun(t, y)`` (optionally ``fun(t, y, Fk)`` tolerated).
        t : float
            Current time of the known state ``y``.
        y : ndarray, shape (n,)
            Current state.
        h : float
            Proposed step size.

        Returns
        -------
        y_new : ndarray
            Next state.
        Fk_new : ndarray or None
            Residual / implicit function evaluation at the converged iterate.
        err_est : float
            Solver's local nonlinear residual norm (not LTE). For multi-stage
            composites this is the last stage residual norm.
        success : bool
            True if nonlinear solve converged to tolerance.
        iterations : int
            Number of nonlinear iterations executed.
        """
        A_local = self._get_A(len(y))

        # When h changes the iteration matrix A/h − J changes, so the
        # cached SPLU from a different step size cannot be reused directly.
        # However, Richardson extrapolation *alternates* between h and h/2
        # every adaptive step.  Instead of discarding the factorisation we
        # save it under the old step size and restore a cached one (if any)
        # for the new step size.  This avoids re-factorising a 33k × 33k
        # system twice per adaptive step when the Jacobian hasn't changed.
        _prev_h = getattr(self, '_cached_step_h', None)
        if _prev_h is not None and _prev_h != h:
            # ---- save current LU state under the old step size ----
            _cache = getattr(self, '_lu_h_cache', None)
            if _cache is None:
                self._lu_h_cache = _cache = {}
            _cache[_prev_h] = (
                self.solver._lu,
                getattr(self.solver, '_lu_shape', None),
                getattr(self.solver, '_lu_pattern', None),
                getattr(self.solver, '_lu_use_count', 0),
                getattr(self.solver, '_J_cross_call', None),
            )
            # keep at most 2 entries (full-step + half-step)
            if len(_cache) > 2:
                _oldest = next(iter(_cache))
                del _cache[_oldest]

            # ---- restore cached state for the new step size ----
            _cached_state = _cache.get(h)
            if _cached_state is not None:
                (self.solver._lu,
                 self.solver._lu_shape,
                 self.solver._lu_pattern,
                 self.solver._lu_use_count,
                 self.solver._J_cross_call) = _cached_state
            else:
                self.solver._lu = None
                self.solver._lu_shape = None
                if hasattr(self.solver, '_J_cross_call'):
                    self.solver._J_cross_call = None
            # PETSc/MUMPS factorisation cannot be swapped like SPLU;
            # force a refactorisation on the next linear solve.
            self.solver._petsc_needs_matrix_update = True
        if _prev_h is None:
            # First call ever — no state to save/restore
            pass
        self._cached_step_h = h

        # Helper to flexibly call fun with optional Fk_val
        prev_state_arg = y if self.pass_prev_state else None
        step_size_arg = h if self.pass_step_size else None

        # Bind RHS once with available context (prev_state, step_size)
        _rhs = self._get_bound_wrapper(fun, has_prev=(prev_state_arg is not None), has_h=(step_size_arg is not None), cache=self._fun_bindings)

        def _call_fun(f, tt, yy, Fk=None):
            # f is ignored; wrapper is bound to func shape
            return _rhs(tt, yy, Fk, prev_state_arg, step_size_arg)

        # Define the implicit equation for the backward Euler step.
        implicit_eq = lambda y_new: A_local @ ((y_new - y) / h) - _call_fun(
            fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None)
        )

        # If an analytical RHS Jacobian is provided, set an exact residual Jacobian for this step.
        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac) and getattr(self.solver, 'method', None) != 'VI':
            A_over_h = A_local / h

            def jac_eq(y_new, _Aoh=A_over_h, _t=t, _h=h, _rhs_jac=rhs_jac, _solver=self.solver):
                # J_res = (A_local/h) - d f/dy (t+h, y_new)
                fk_val = getattr(_solver, 'last_Fk_val', None)
                _jac = self._get_bound_wrapper(_rhs_jac, has_prev=(prev_state_arg is not None), has_h=(step_size_arg is not None), cache=self._jac_bindings)
                rhs_jac_val = _jac(_t + _h, y_new, fk_val, prev_state_arg, step_size_arg)
                return _Aoh - rhs_jac_val

            self.solver.jacobian = jac_eq

        # Thread step-context to the solver so projections can access it
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = h
        except Exception:
            pass

        return self.solver.solve(implicit_eq, y)


class AlgebraicBackwardEuler(IntegrationMethod):
    """Backward Euler variant supporting a subset of algebraic (index-1) constraints.

    Differential indices D use standard BE residual:
        (y_new - y)/h - f(t+h, y_new) = 0   (or A-scaled if A provided)

    Algebraic indices A enforce g(y_new)_A = 0 directly (no time derivative term).

    Parameters
    ----------
    algebraic_indices : sequence[int]
        Indices in the state treated as algebraic.
    g_func : callable
        Constraint function g(y) returning vector with at least entries for algebraic indices.
        Only entries at algebraic indices are used. If None, will default to identity constraint
        (i.e. forces y_A = 0).
    g_jac : callable, optional
        Jacobian dg/dy. If provided, an exact residual Jacobian is formed for algebraic rows.
    solver : ImplicitEquationSolver, optional
        Nonlinear solver to use (must be semismooth_newton or VI with projection).
    A : ndarray or None
        Mass / coefficient matrix for differential part; identity if None.
    """
    def __init__(self, algebraic_indices, g_func=None, g_jac=None, solver=None, A=None):
        self.algebraic_indices = np.array(sorted(algebraic_indices), dtype=int)
        self.g_func = g_func
        self.g_jac = g_jac
        self.solver = solver or ImplicitEquationSolver(method='semismooth_newton')
        self.A = A
        self.use_identity = A is None

    def _get_A(self, n):
        if not self.use_identity:
            return self.A
        # Prefer CSR identity for large n to reduce memory/compute
        want_sparse = False
        try:
            want_sparse = bool(self.solver._sparse_active(n))
        except Exception:
            thr = getattr(self.solver, 'sparse_threshold', 200)
            try:
                want_sparse = (n >= int(thr))
            except Exception:
                want_sparse = (n >= 200)
        if want_sparse:
            return sp.eye(n, format='csr')
        return np.eye(n)

    def step(self, fun, t, y, h):
        n = len(y)
        A_local = self._get_A(n)
        alg_idx = self.algebraic_indices
        diff_mask = np.ones(n, dtype=bool)
        diff_mask[alg_idx] = False
        diff_idx = np.nonzero(diff_mask)[0]

        def _call_fun(f, tt, yy, Fk=None):
            try:
                return f(tt, yy, Fk)
            except TypeError:
                try:
                    return f(tt, yy)
                except TypeError:
                    return f(yy)

        def residual(y_new):
            F = np.zeros(n, dtype=float)
            f_val = _call_fun(fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
            # Differential part rows
            if diff_idx.size:
                F[diff_idx] = (A_local[diff_idx][:, :] @ ((y_new - y) / h)) - f_val[diff_idx]
            # Algebraic part rows: g(y_new)_A = 0
            if alg_idx.size:
                if self.g_func is None:
                    F[alg_idx] = y_new[alg_idx]  # default constraint y_A = 0
                else:
                    g_val = self.g_func(y_new)
                    # Allow g to return full vector or just algebraic subset
                    if g_val.shape[0] == n:
                        F[alg_idx] = g_val[alg_idx]
                    else:
                        # assume ordered alignment with alg_idx
                        F[alg_idx] = g_val
            return F

        # Attach Jacobian if possible
        def jacobian(y_new):
            J = np.zeros((n, n), dtype=float)
            # Differential rows
            f_val = _call_fun(fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
            # Need df/dy for diff rows; reuse solver.rhs_jacobian if present
            rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
            if callable(rhs_jac):
                try:
                    dfdy = rhs_jac(t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
                except TypeError:
                    dfdy = rhs_jac(t + h, y_new)
            else:
                # fallback finite-difference (small system expected)
                eps = 1e-8
                dfdy = np.zeros((n, n))
                f0 = _call_fun(fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
                for j in range(n):
                    y_pert = y_new.copy(); y_pert[j] += eps
                    f_eps = _call_fun(fun, t + h, y_pert, getattr(self.solver, 'last_Fk_val', None))
                    dfdy[:, j] = (f_eps - f0)/eps
            if diff_idx.size:
                J[diff_idx, :] = (A_local[diff_idx][:, :] / h) - dfdy[diff_idx, :]
            if alg_idx.size:
                if self.g_jac is not None:
                    gJ = self.g_jac(y_new)
                    if gJ.shape[0] == len(alg_idx):
                        for row_pos, gi in enumerate(alg_idx):
                            J[gi, :] = gJ[row_pos]
                    else:
                        # assume full sized
                        J[alg_idx, :] = gJ[alg_idx, :]
                else:
                    # Numerical for g only
                    if self.g_func is None:
                        J[alg_idx, alg_idx] = 1.0
                    else:
                        eps = 1e-8
                        g0 = self.g_func(y_new)
                        full_g0 = np.zeros(len(alg_idx)) if g0.shape[0]==n else g0
                        if g0.shape[0]==n:
                            full_g0 = g0[alg_idx]
                        for j in range(n):
                            y_pert = y_new.copy(); y_pert[j] += eps
                            g_eps = self.g_func(y_pert)
                            if g_eps.shape[0]==n:
                                g_eps_sub = g_eps[alg_idx]
                            else:
                                g_eps_sub = g_eps
                            J[alg_idx, j] = (g_eps_sub - full_g0)/eps
            return J

        # Provide Jacobian to solver (only if semismooth_newton path)
        if getattr(self.solver, 'method', None) != 'VI':
            self.solver.jacobian = jacobian

        # Thread step-context
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = h
        except Exception:
            pass

        return self.solver.solve(residual, y)


class Trapezoidal(BackwardEuler):
    def __init__(self, solver=None, A=None, **kwargs):
        super().__init__(solver=solver, A=A, **kwargs)
        self.order = 2
    """
    Implements the Trapezoidal (Crank-Nicolson) integration method.

    Inherits the matrix handling from BackwardEuler.
    """

    def step(self, fun, t, y, h):
        """Perform one implicit Trapezoidal (Crank–Nicolson) step.

        Nonlinear system::

            A ((y_new - y)/h) - 0.5 ( f(t, y) + f(t+h, y_new) ) = 0

        Returns follow the same 5-tuple convention described in
        :meth:`BackwardEuler.step`.
        """
        A_local = self._get_A(len(y))

        prev_state_arg = y if self.pass_prev_state else None
        step_size_arg = h if self.pass_step_size else None

        # Bind RHS once (allows step_size override at call)
        _rhs = self._get_bound_wrapper(fun, has_prev=(prev_state_arg is not None), has_h=True, cache=self._fun_bindings)

        def _call_fun(f, tt, yy, Fk=None, h_override=None):
            return _rhs(tt, yy, Fk, prev_state_arg, (h_override if h_override is not None else step_size_arg))

        fk_last = getattr(self.solver, 'last_Fk_val', None)
        f_n = _call_fun(fun, t, y, fk_last)
        implicit_eq = lambda y_new: A_local @ ((y_new - y) / h) - 0.5 * (
            f_n + _call_fun(fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
        )

        # Exact Jacobian when analytical RHS Jacobian is available
        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac) and getattr(self.solver, 'method', None) != 'VI':
            A_over_h = A_local / h

            def jac_eq(y_new, _Aoh=A_over_h, _t=t, _h=h, _rhs_jac=rhs_jac, _solver=self.solver):
                # J_res = (A_local/h) - 0.5 * d f/dy (t+h, y_new)
                fk_val = getattr(_solver, 'last_Fk_val', None)
                _jac = self._get_bound_wrapper(_rhs_jac, has_prev=(prev_state_arg is not None), has_h=(step_size_arg is not None), cache=self._jac_bindings)
                rhs_jac_val = _jac(_t + _h, y_new, fk_val, prev_state_arg, step_size_arg)
                return _Aoh - 0.5 * rhs_jac_val

            self.solver.jacobian = jac_eq

        # Thread step-context
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = h
        except Exception:
            pass

        return self.solver.solve(implicit_eq, y)


class ThetaMethod(BackwardEuler):
    """
    Implements the Theta integration method, a generalization of Backward Euler and Trapezoidal methods.

    The method uses a parameter theta in [0, 1]:
      - theta = 1 gives Backward Euler,
      - theta = 0.5 gives Trapezoidal method.
    """

    def __init__(self, theta=0.5, **kwargs):
        """
        Initialize a ThetaMethod instance.

        Parameters:
            theta: float
                The weighting parameter between 0 and 1.
            **kwargs:
                Additional keyword arguments passed to the BackwardEuler initializer.

        Raises:
            ValueError: if theta is not in the interval [0, 1].
        """
        super().__init__(**kwargs)
        if not (0 <= theta <= 1):
            raise ValueError("Theta must be between 0 and 1")
        self.theta = theta
        # Theta=0.5 is TR (order 2); otherwise default to order 1
        self.order = 2 if abs(self.theta - 0.5) < 1e-12 else 1

    def step(self, fun, t, y, h):
        """Perform one Theta method step.

        Nonlinear system::

            A ((y_new - y)/h) - ( theta f(t+h, y_new) + (1-\theta) f(t, y) ) = 0

        Returns: 5-tuple as in :meth:`BackwardEuler.step`.
        """
        A_local = self._get_A(len(y))

        prev_state_arg = y if self.pass_prev_state else None
        step_size_arg = h if self.pass_step_size else None

        _rhs = self._get_bound_wrapper(fun, has_prev=(prev_state_arg is not None), has_h=(step_size_arg is not None), cache=self._fun_bindings)

        def _call_fun(f, tt, yy, Fk=None):
            return _rhs(tt, yy, Fk, prev_state_arg, step_size_arg)

        fk_last = getattr(self.solver, 'last_Fk_val', None)
        f_n = _call_fun(fun, t, y, fk_last)
        implicit_eq = lambda y_new: A_local @ ((y_new - y) / h) - (
            self.theta * _call_fun(fun, t+h, y_new, getattr(self.solver, 'last_Fk_val', None)) + (1 - self.theta) * f_n
        )

        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac) and getattr(self.solver, 'method', None) != 'VI':
            A_over_h = A_local / h
            theta_val = self.theta

            def jac_eq(y_new, _Aoh=A_over_h, _t=t, _h=h, _rhs_jac=rhs_jac, _theta=theta_val, _solver=self.solver):
                # J_res = (A_local/h) - theta * d f/dy (t+h, y_new)
                fk_val = getattr(_solver, 'last_Fk_val', None)
                _jac = self._get_bound_wrapper(_rhs_jac, has_prev=(prev_state_arg is not None), has_h=(step_size_arg is not None), cache=self._jac_bindings)
                rhs_jac_val = _jac(_t + _h, y_new, fk_val, prev_state_arg, step_size_arg)
                return _Aoh - _theta * rhs_jac_val

            self.solver.jacobian = jac_eq

        # Thread step-context
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = h
        except Exception:
            pass

        return self.solver.solve(implicit_eq, y)


class CompositeMethod(IntegrationMethod):
    """
    Implements a composite integration method that combines two steps:
      1. A half-step using the Trapezoidal method.
      2. A full step using a modified Backward Euler method.

    The composite method first advances the solution halfway in time, then uses this intermediate
    value to compute the final step.
    """

    def __init__(self, a=1.0, solver=None, A=None, **kwargs):
        """
        Initialize the CompositeMethod.

        Parameters:
            a: float, optional
                A parameter that may be used for weighting (currently not used in the implementation).
            solver: ImplicitEquationSolver, optional
                The solver used to solve the implicit equations.
            A: np.array, optional
                The matrix used in the formulation. If None, identity is used.
        """
        self.solver = solver or ImplicitEquationSolver(method='semismooth_newton')
        # Create instances for sub-steps using Trapezoidal and Backward Euler methods.
        self.trapezoidal = Trapezoidal(solver=self.solver, A=A, **kwargs)
        self.backward_euler = BackwardEuler(solver=self.solver, A=A, **kwargs)
        # TR-BE composite is second-order (TR-BDF2 style)
        self.order = 2

    def step(self, fun, t, y, h):
        """Composite TR / BE second-order step (TR-BDF2 style variant).

        Two stages:
          1. Half-step TR to obtain ``y_half``.
          2. Modified BE relation ``(3*y_new - 4*y_half + y)/h - f(t+h, y_new) = 0``.

        Returns
        -------
        y_new, Fk_new, err_new, success, iterations : as in other integrators with
        iteration count the sum over both stages.
        """
        half_h = 0.5 * h

        # ----- Stage 1: half-step TR from (t, y) to (t+half_h, y_half)
        # The TR.step call itself will set prev_state=y etc. via shared solver.
        y_half, Fk_half, err_half, success_half, iters_half = \
            self.trapezoidal.step(fun, t, y, half_h)
        if not success_half:
            return y, Fk_half, err_half, False, iters_half

        # ----- Stage 2: BE-like relation from (t+half_h, y_half) to (t+h, y_new)
        prev_state_arg = y_half if getattr(self.backward_euler, 'pass_prev_state', False) else None
        step_size_arg = h if getattr(self.backward_euler, 'pass_step_size', False) else None

        # Bind RHS once for the second stage using BackwardEuler's binder/caches
        _rhs_be = self.backward_euler._get_bound_wrapper(
            fun,
            has_prev=(prev_state_arg is not None),
            has_h=(step_size_arg is not None),
            cache=self.backward_euler._fun_bindings,
        )

        def _call_fun(f, tt, yy, Fk=None):
            return _rhs_be(tt, yy, Fk, prev_state_arg, step_size_arg)

        def implicit_eq(y_new):
            A_local = self.backward_euler._get_A(len(y))
            return A_local @ ((3.0 * y_new - 4.0 * y_half + y) / h) - _call_fun(
                fun, t + h, y_new, getattr(self.backward_euler.solver, 'last_Fk_val', None)
            )

        rhs_jac = getattr(self.backward_euler.solver, 'rhs_jacobian', None)
        if callable(rhs_jac) and getattr(self.backward_euler.solver, 'method', None) != 'VI':
            A_over_h = self.backward_euler._get_A(len(y)) / h

            def jac_eq_second(y_new, _Aoh=A_over_h, _t=t, _h=h, _rhs_jac=rhs_jac, _solver=self.backward_euler.solver):
                fk_val = getattr(_solver, 'last_Fk_val', None)
                _jac = self.backward_euler._get_bound_wrapper(
                    _rhs_jac,
                    has_prev=(prev_state_arg is not None),
                    has_h=(step_size_arg is not None),
                    cache=self.backward_euler._jac_bindings,
                )
                rhs_jac_val = _jac(_t + _h, y_new, fk_val, prev_state_arg, step_size_arg)
                return (3.0 * _Aoh) - rhs_jac_val

            self.backward_euler.solver.jacobian = jac_eq_second

        # Thread step-context for stage 2 (previous accepted state is y_half)
        try:
            self.backward_euler.solver.current_time = t + h
            self.backward_euler.solver.prev_state = y_half
            self.backward_euler.solver.prev_time = t + half_h
            self.backward_euler.solver.prev_step = half_h
        except Exception:
            pass

        y_guess = y_half
        y_new, Fk_new, err_new, success_new, iters_new = self.backward_euler.solver.solve(implicit_eq, y_guess)
        total_iters = iters_half + iters_new
        overall_success = success_half and success_new
        return (y_new, Fk_new, err_new, overall_success, total_iters)


class SDIRK2(BackwardEuler):
    r"""Two-stage, L-stable, singly diagonally implicit Runge–Kutta method of
    order 2 with an embedded order-1 error estimate.

    Butcher tableau (Alexander 1977)::

        γ   | γ     0
        1   | 1-γ   γ
        ----|----------
            | 1-γ   γ     ← order 2
            | 1     0     ← order 1 (embedded)

    where γ = 1 − √2/2 ≈ 0.2929.

    Each stage requires one implicit solve with the *same* diagonal
    coefficient γ, which means the iteration matrix
    ``(A/(γh) − df/dy)`` can in principle be reused across both stages.

    The embedded pair yields a cheap local-truncation-error estimate
    expressed purely in terms of the stage values (mass-matrix agnostic)::

        err = (Y₂ − y_shift) − (Y₁ − y)

    where ``y_shift = y + ((1−γ)/γ)(Y₁ − y)``.  This is returned as
    the ``err`` component of the 5-tuple.  The adaptive stepper
    (:class:`AdaptiveStepping`) can therefore skip the three-solve
    Richardson extrapolation when this integrator is used.

    Parameters
    ----------
    solver : ImplicitEquationSolver, optional
        Nonlinear solver instance.
    A : ndarray or sparse matrix, optional
        Mass / descriptor matrix; identity if *None*.
    pass_prev_state, pass_step_size : bool
        Forwarded to :class:`BackwardEuler`.
    """

    # Diagonal coefficient  γ = 1 − √2/2
    _GAMMA = 1.0 - math.sqrt(2.0) / 2.0

    def __init__(self, solver=None, A=None, **kwargs):
        super().__init__(solver=solver, A=A, **kwargs)
        self.order = 2  # method order for the adaptive controller
        self.has_embedded_error = True  # bypass Richardson in AdaptiveStepping

    # ------------------------------------------------------------------
    def step(self, fun, t, y, h):
        r"""Advance by one SDIRK2 step of size *h*.

        Returns
        -------
        y_new : ndarray
            State at ``t + h`` (second-order accurate).
        Fk_new : ndarray or None
            Residual at the converged iterate of stage 2.
        err : ndarray
            Element-wise embedded error estimate ``(Y₂ − y_shift) − (Y₁ − y)``,
            equivalent to ``h γ (K₂ − K₁)`` where ``K_i`` are the stage
            derivatives.  Mass-matrix agnostic.
        success : bool
            True if *both* stage solves converged.
        iterations : int
            Sum of nonlinear iterations over the two stages.
        """
        gamma = self._GAMMA
        n = len(y)
        A_local = self._get_A(n)
        gh = gamma * h  # diagonal sub-step length

        # -- Cache A/(γh) across steps to avoid redundant sparse division --
        # Both SDIRK2 stages share the same diagonal coefficient γh, so the
        # matrix  A/(γh)  only changes when h changes.  Caching it avoids a
        # full sparse-scalar division (O(nnz) allocation + copy) on every
        # call to step() when h is unchanged — which is the common case for
        # an adaptive controller on a smooth trajectory.
        _prev_gh = getattr(self, '_cached_gh', None)
        if _prev_gh is None or _prev_gh != gh:
            self._A_over_gh = A_local / gh
            self._cached_gh = gh
            # ---- Invalidate solver caches ----
            # The iteration matrix  M/(γh) − J  depends on h.  When h
            # changes, a cached SPLU factorisation or Jacobian from the
            # previous step is stale and MUST be discarded.  Without
            # this, the modified-Newton loop's first iteration reuses
            # the stale factor (because errF > 0.5*inf is always False
            # in IEEE-754), corrupts the iterate, and wastes iterations
            # recovering — often exceeding max_iter.
            self.solver._lu = None
            self.solver._lu_shape = None
            if hasattr(self.solver, '_J_cross_call'):
                self.solver._J_cross_call = None
            # Also invalidate PETSc/MUMPS factorisation — the direct
            # solver path in _solve_with_petsc will otherwise skip the
            # matrix update and keep solving with the stale factor.
            self.solver._petsc_needs_matrix_update = True
        A_over_gh = self._A_over_gh

        # -- helpers for flexible RHS calling ----------------------------
        prev_state_arg = y if self.pass_prev_state else None
        # Pass the *stage* sub-step γh (not full h).  The contact RHS
        # computes  B·r/(θ·h_val);  the stage equation divides by γh,
        # so net coupling is  B·r·γh / (θ·γh) = B·r/θ — correct for
        # θ = 1.  Pre-stress callbacks (get_s0) that multiply force×h_val
        # also get the stage-level impulse, which is physically right.
        step_size_arg = gh if self.pass_step_size else None
        _rhs = self._get_bound_wrapper(
            fun,
            has_prev=(prev_state_arg is not None),
            has_h=(step_size_arg is not None),
            cache=self._fun_bindings,
        )

        def _call_fun(tt, yy, Fk=None):
            return _rhs(tt, yy, Fk, prev_state_arg, step_size_arg)

        # ── Block-diagonal ρ equilibration ──────────────────────────────
        # The natural-map Jacobian  J_Φ = I − DΠ(I − ρ J_F)  has the
        # implicit-equation diagonal  A_ii/(γh)  on the physical rows.
        # When γ < 1 (e.g. SDIRK2, γ ≈ 0.293) this diagonal scales
        # as 1/γ relative to Backward Euler, pushing the condition
        # number above the IEEE-754 accuracy floor and stalling Newton.
        # Setting  ρ_i = γh / A_ii  for M-rows and  ρ_i = 1  for
        # algebraic (zero-mass) rows equilibrates the system to O(1).
        _A_diag = (A_local.diagonal() if sp.issparse(A_local)
                   else np.diag(A_local))
        _phys_mask = np.abs(_A_diag) > 0.0
        _rho_vec = np.ones(n)
        _rho_vec[_phys_mask] = gh / _A_diag[_phys_mask]
        self.solver.lam = _rho_vec

        # ================================================================
        # Stage 1:  solve  A (Y1 − y)/(γh) − f(t + γh, Y1) = 0
        # ================================================================
        def implicit_s1(Y1):
            return A_over_gh @ (Y1 - y) - _call_fun(
                t + gh, Y1, getattr(self.solver, 'last_Fk_val', None)
            )

        # Exact Jacobian for stage 1 (if available)
        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac) and getattr(self.solver, 'method', None) != 'VI':

            def jac_s1(Y1, _Aogh=A_over_gh, _t=t, _gh=gh):
                fk_val = getattr(self.solver, 'last_Fk_val', None)
                _jac = self._get_bound_wrapper(
                    rhs_jac,
                    has_prev=(prev_state_arg is not None),
                    has_h=(step_size_arg is not None),
                    cache=self._jac_bindings,
                )
                J_rhs = _jac(_t + _gh, Y1, fk_val, prev_state_arg, step_size_arg)
                return _Aogh - J_rhs

            self.solver.jacobian = jac_s1

        # Thread step-context for stage 1
        try:
            self.solver.current_time = t + gh
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = gh
        except Exception:
            pass

        Y1, Fk1, err1, ok1, it1 = self.solver.solve(implicit_s1, y)
        if not ok1:
            return y, Fk1, np.zeros(n), False, it1

        # ================================================================
        # Stage 2:  solve  A (Y2 − y_shift) / (γh) − f(t+h, Y2) = 0
        #
        # The Butcher stage derivative K1 = (Y1 − y)/(γh).  The stage-2
        # shift uses K1 directly (no f evaluation or M⁻¹ needed):
        #   y_shift = y + (1−γ)h K1 = y + ((1−γ)/γ)(Y1 − y)
        # ================================================================
        dY1 = Y1 - y                                   # stage-1 increment
        y_shift = y + ((1.0 - gamma) / gamma) * dY1    # mass-matrix agnostic

        def implicit_s2(Y2):
            return A_over_gh @ (Y2 - y_shift) - _call_fun(
                t + h, Y2, getattr(self.solver, 'last_Fk_val', None)
            )

        # Exact Jacobian for stage 2 (same structure, different time)
        if callable(rhs_jac) and getattr(self.solver, 'method', None) != 'VI':

            def jac_s2(Y2, _Aogh=A_over_gh, _t=t, _h=h):
                fk_val = getattr(self.solver, 'last_Fk_val', None)
                _jac = self._get_bound_wrapper(
                    rhs_jac,
                    has_prev=(prev_state_arg is not None),
                    has_h=(step_size_arg is not None),
                    cache=self._jac_bindings,
                )
                J_rhs = _jac(_t + _h, Y2, fk_val, prev_state_arg, step_size_arg)
                return _Aogh - J_rhs

            self.solver.jacobian = jac_s2

        # Thread step-context for stage 2
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = Y1
            self.solver.prev_time = t + gh
            self.solver.prev_step = gh
        except Exception:
            pass

        # Good initial guess: extrapolate from stage 1
        Y2_guess = y_shift
        Y2, Fk2, err2, ok2, it2 = self.solver.solve(implicit_s2, Y2_guess)
        if not ok2:
            return y, Fk2, np.zeros(n), False, it1 + it2

        # ================================================================
        # Output
        # ================================================================
        # y_{n+1} = Y2  (already the second-order solution)
        y_new = Y2

        # Embedded error estimate (mass-matrix agnostic, uses stage values):
        #   err = h γ (K2 − K1)
        #       = (Y2 − y_shift) − (Y1 − y)
        #       = (Y2 − y_shift) − dY1
        err_embed = (Y2 - y_shift) - dY1

        total_iters = it1 + it2
        return y_new, Fk2, err_embed, True, total_iters


class EmbeddedBETR(IntegrationMethod):
    """
    Trapezoidal implicit integrator (kept under the historical name EmbeddedBETR).

    Acts as a plain TR stepper with optional exact residual Jacobian.
    """

    def __init__(self, solver=None, A=None):
        self.solver = solver or ImplicitEquationSolver(method='semismooth_newton')
        self.A = A
        self.use_identity = (A is None)
        self.order = 2

    def _get_A(self, n):
        if not self.use_identity:
            return self.A
        want_sparse = False
        try:
            want_sparse = bool(self.solver._sparse_active(n))
        except Exception:
            thr = getattr(self.solver, 'sparse_threshold', 200)
            try:
                want_sparse = (n >= int(thr))
            except Exception:
                want_sparse = (n >= 200)
        if want_sparse:
            return sp.eye(n, format='csr')
        return np.eye(n)

    def _attach_be_jac(self, A_local, t, h):
        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac):
            def jac_eq(y_new, _A=A_local, _h=h, _t=t, _rhs_jac=rhs_jac):
                return (_A / _h) - _rhs_jac(_t + _h, y_new)
            self.solver.jacobian = jac_eq

    def _attach_tr_jac(self, A_local, t, h):
        rhs_jac = getattr(self.solver, 'rhs_jacobian', None)
        if callable(rhs_jac):
            def jac_eq(y_new, _A=A_local, _h=h, _t=t, _rhs_jac=rhs_jac, _solver=self.solver):
                try:
                    return (_A / _h) - 0.5 * _rhs_jac(_t + _h, y_new, getattr(_solver, 'last_Fk_val', None))
                except TypeError:
                    return (_A / _h) - 0.5 * _rhs_jac(_t + _h, y_new)
            self.solver.jacobian = jac_eq

    def step(self, fun, t, y, h):
        """Single TR step (legacy name kept for backward compatibility).

        Returns
        -------
        y_new, Fk_new, err_new, success, iterations : standard solver 5-tuple.
        """
        n = len(y)
        A_local = self._get_A(n)

        def _call_fun(f, tt, yy, Fk=None):
            try:
                return f(tt, yy, Fk)
            except TypeError:
                try:
                    return f(tt, yy)
                except TypeError:
                    return f(yy)

        f_n = _call_fun(fun, t, y, getattr(self.solver, 'last_Fk_val', None))
        implicit_tr = lambda y_new: A_local @ ((y_new - y) / h) - 0.5 * (
            f_n + _call_fun(fun, t + h, y_new, getattr(self.solver, 'last_Fk_val', None))
        )
        # Attach TR residual Jacobian using shared helper
        self._attach_tr_jac(A_local, t, h)

        # Thread step-context
        try:
            self.solver.current_time = t + h
            self.solver.prev_state = y
            self.solver.prev_time = t
            self.solver.prev_step = h
        except Exception:
            pass

        return self.solver.solve(implicit_tr, y)


if __name__ == "__main__":
    # Optional quick smoke test when running this module directly
    def rhs(t, y, Fk=None):
        return -y

    y0 = np.array([1.0])
    t0 = 0.0
    h = 0.2
    be = BackwardEuler(solver=ImplicitEquationSolver(method='semismooth_newton', proj=lambda *args, **kw: None))
    tr = Trapezoidal(solver=ImplicitEquationSolver(method='semismooth_newton', proj=lambda *args, **kw: None))
    comp = CompositeMethod(solver=ImplicitEquationSolver(method='semismooth_newton', proj=lambda *args, **kw: None))
    print("TR:", tr.step(rhs, t0, y0.copy(), h))
    print("Composite:", comp.step(rhs, t0, y0.copy(), h))
