from collections.abc import Callable, Sequence
from typing import Dict, Tuple, Optional
from functools import partial

import keras

import numpy as np
from typing import Literal, Union

from bayesflow.adapters import Adapter
from bayesflow.types import Tensor
from bayesflow.utils import filter_kwargs
from bayesflow.utils.logging import warning, debug


ArrayLike = int | float | Tensor
StateDict = Dict[str, ArrayLike]


DETERMINISTIC_METHODS = ["euler", "rk45", "tsit5"]
STOCHASTIC_METHODS = ["euler_maruyama", "sea", "shark", "two_step_adaptive", "langevin"]


def _check_all_nans(state: StateDict):
    all_nans_flags = []
    for v in state.values():
        all_nans_flags.append(keras.ops.all(keras.ops.isnan(v)))
    return keras.ops.all(keras.ops.stack(all_nans_flags))


def euler_step(
    fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
) -> Tuple[StateDict, ArrayLike, None, ArrayLike]:
    k1 = fn(time, **filter_kwargs(state, fn))

    new_state = state.copy()
    for key in k1.keys():
        new_state[key] = state[key] + step_size * k1[key]
    new_time = time + step_size

    return new_state, new_time, None, 0.0


def add_scaled(state, ks, coeffs, h):
    out = {}
    for key, y in state.items():
        acc = keras.ops.zeros_like(y)
        for c, k in zip(coeffs, ks):
            acc = acc + c * k[key]
        out[key] = y + h * acc
    return out


def rk45_step(
    fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    k1: StateDict = None,
    use_adaptive_step_size: bool = True,
) -> Tuple[StateDict, ArrayLike, StateDict | None, ArrayLike]:
    """
    Dormand-Prince 5(4) method with embedded error estimation [1].

    Dormand (1996), Numerical Methods for Differential Equations: A Computational Approach
    """
    h = step_size

    if k1 is None:  # reuse k1 if available
        k1 = fn(time, **filter_kwargs(state, fn))
    k2 = fn(time + h * (1 / 5), **filter_kwargs(add_scaled(state, [k1], [1 / 5], h), fn))
    k3 = fn(time + h * (3 / 10), **filter_kwargs(add_scaled(state, [k1, k2], [3 / 40, 9 / 40], h), fn))
    k4 = fn(time + h * (4 / 5), **filter_kwargs(add_scaled(state, [k1, k2, k3], [44 / 45, -56 / 15, 32 / 9], h), fn))
    k5 = fn(
        time + h * (8 / 9),
        **filter_kwargs(
            add_scaled(state, [k1, k2, k3, k4], [19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729], h), fn
        ),
    )
    k6 = fn(
        time + h,
        **filter_kwargs(
            add_scaled(state, [k1, k2, k3, k4, k5], [9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656], h),
            fn,
        ),
    )

    # 5th order solution
    new_state = {}
    for key in k1.keys():
        new_state[key] = state[key] + h * (
            35 / 384 * k1[key] + 500 / 1113 * k3[key] + 125 / 192 * k4[key] - 2187 / 6784 * k5[key] + 11 / 84 * k6[key]
        )

    new_time = time + h
    if not use_adaptive_step_size:
        return new_state, new_time, None, 0.0

    k7 = fn(time + h, **filter_kwargs(new_state, fn))

    # 4th order embedded solution
    err_state = {}
    for key in k1.keys():
        y4 = state[key] + h * (
            5179 / 57600 * k1[key]
            + 7571 / 16695 * k3[key]
            + 393 / 640 * k4[key]
            - 92097 / 339200 * k5[key]
            + 187 / 2100 * k6[key]
            + 1 / 40 * k7[key]
        )
        err_state[key] = new_state[key] - y4

    err_norm = keras.ops.stack([keras.ops.norm(v, ord=2, axis=-1) for v in err_state.values()])
    err = keras.ops.max(err_norm)

    return new_state, new_time, k7, err


def tsit5_step(
    fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    k1: StateDict = None,
    use_adaptive_step_size: bool = True,
) -> Tuple[StateDict, ArrayLike, StateDict | None, ArrayLike]:
    """
    Implements a single step of the Tsitouras 5/4 Runge-Kutta method [1].

    [1] Tsitouras (2011), Runge--Kutta pairs of order 5(4) satisfying only the first column simplifying assumption
    """
    h = step_size

    # Butcher tableau coefficients
    c2 = 0.161
    c3 = 0.327
    c4 = 0.9
    c5 = 0.9800255409045097

    if k1 is None:  # reuse k1 if available
        k1 = fn(time, **filter_kwargs(state, fn))
    k2 = fn(time + h * c2, **filter_kwargs(add_scaled(state, [k1], [0.161], h), fn))
    k3 = fn(
        time + h * c3, **filter_kwargs(add_scaled(state, [k1, k2], [-0.0084806554923570, 0.3354806554923570], h), fn)
    )
    k4 = fn(
        time + h * c4,
        **filter_kwargs(
            add_scaled(state, [k1, k2, k3], [2.897153057105494, -6.359448489975075, 4.362295432869581], h), fn
        ),
    )
    k5 = fn(
        time + h * c5,
        **filter_kwargs(
            add_scaled(
                state,
                [k1, k2, k3, k4],
                [5.325864828439257, -11.74888356406283, 7.495539342889836, -0.09249506636175525],
                h,
            ),
            fn,
        ),
    )
    k6 = fn(
        time + h,
        **filter_kwargs(
            add_scaled(
                state,
                [k1, k2, k3, k4, k5],
                [5.86145544294270, -12.92096931784711, 8.159367898576159, -0.07158497328140100, -0.02826905039406838],
                h,
            ),
            fn,
        ),
    )

    # 5th order solution: b coefficients
    new_state = {}
    for key in state.keys():
        new_state[key] = state[key] + h * (
            0.09646076681806523 * k1[key]
            + 0.01 * k2[key]
            + 0.4798896504144996 * k3[key]
            + 1.379008574103742 * k4[key]
            - 3.290069515436081 * k5[key]
            + 2.324710524099774 * k6[key]
        )

    new_time = time + h
    if not use_adaptive_step_size:
        return new_state, new_time, None, 0.0

    k7 = fn(time + h, **filter_kwargs(new_state, fn))

    err_state = {}
    for key in state.keys():
        err_state[key] = h * (
            0.007880878010261995 * k3[key]
            + 0.5823571654525552 * k5[key]
            + 0.015151515151515152 * k7[key]
            - 0.00178001105222577714 * k1[key]
            - 0.0008164344596567469 * k2[key]
            - 0.1447110071732629 * k4[key]
            - 0.45808210592918697 * k6[key]
        )

    err_norm = keras.ops.stack([keras.ops.norm(v, ord=2, axis=-1) for v in err_state.values()])
    err = keras.ops.max(err_norm)

    return new_state, new_time, k7, err


def integrate_fixed(
    fn: Callable,
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    steps: int,
    method: str,
    **kwargs,
) -> StateDict:
    if steps <= 0:
        raise ValueError("Number of steps must be positive.")

    match method:
        case "euler":
            step_fn = partial(euler_step, fn, **filter_kwargs(kwargs, rk45_step))
        case "rk45":
            step_fn = partial(rk45_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=False)
        case "tsit5":
            step_fn = partial(tsit5_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=False)
        case str() as name:
            raise ValueError(f"Unknown integration method name: {name!r}")
        case other:
            raise TypeError(f"Invalid integration method: {other!r}")

    step_size = (stop_time - start_time) / steps

    def body(_loop_var, _loop_state):
        _state, _time = _loop_state
        _state, _time, _, _ = step_fn(_state, _time, step_size)
        return _state, _time

    state, _ = keras.ops.fori_loop(
        0,
        steps,
        body,
        (state, start_time),
    )
    return state


def integrate_scheduled(
    fn: Callable,
    state: StateDict,
    steps: Tensor | np.ndarray,
    method: str,
    **kwargs,
) -> StateDict:
    match method:
        case "euler":
            step_fn = partial(euler_step, fn, **filter_kwargs(kwargs, rk45_step))
        case "rk45":
            step_fn = partial(rk45_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=False)
        case "tsit5":
            step_fn = partial(tsit5_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=False)
        case str() as name:
            raise ValueError(f"Unknown integration method name: {name!r}")
        case other:
            raise TypeError(f"Invalid integration method: {other!r}")

    def body(_loop_var, _loop_state):
        _time = steps[_loop_var]
        step_size = steps[_loop_var + 1] - steps[_loop_var]
        _loop_state, _, _, _ = step_fn(_loop_state, _time, step_size)
        return _loop_state

    state = keras.ops.fori_loop(
        0,
        keras.ops.shape(steps)[0] - 1,
        body,
        state,
    )
    return state


def integrate_adaptive(
    fn: Callable,
    state: dict[str, ArrayLike],
    start_time: ArrayLike,
    stop_time: ArrayLike,
    min_steps: int,
    max_steps: int,
    method: str,
    **kwargs,
) -> dict[str, ArrayLike]:
    if max_steps <= min_steps:
        raise ValueError("Maximum number of steps must be greater than minimum number of steps.")

    match method:
        case "rk45":
            step_fn = partial(rk45_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=True)
        case "tsit5":
            step_fn = partial(tsit5_step, fn, **filter_kwargs(kwargs, rk45_step), use_adaptive_step_size=True)
        case "euler":
            raise ValueError("Adaptive step sizing is not supported for the 'euler' method.")
        case str() as name:
            raise ValueError(f"Unknown integration method name: {name!r}")
        case other:
            raise TypeError(f"Invalid integration method: {other!r}")

    atol = keras.ops.convert_to_tensor(kwargs.get("atol", 1e-6), dtype="float32")
    rtol = keras.ops.convert_to_tensor(kwargs.get("rtol", 1e-4), dtype="float32")
    initial_step = keras.ops.convert_to_tensor((stop_time - start_time) / float(min_steps), dtype="float32")
    step0 = keras.ops.convert_to_tensor(0.0, dtype="float32")
    count_not_accepted = 0

    # "First Same As Last" (FSAL) property
    k1_0 = fn(start_time, **filter_kwargs(state, fn))

    def cond(_state, _time, _step_size, _step, _k1, _count_not_accepted):
        time_remaining = keras.ops.sign(stop_time - start_time) * (stop_time - (_time + _step_size))
        step_lt_min = keras.ops.less(_step, float(min_steps))
        step_lt_max = keras.ops.less(_step, float(max_steps))

        all_nans = _check_all_nans(_state)

        end_now = keras.ops.logical_or(
            step_lt_min, keras.ops.logical_and(keras.ops.all(time_remaining > 0), step_lt_max)
        )
        return keras.ops.logical_and(~all_nans, end_now)

    def body(_state, _time, _step_size, _step, _k1, _count_not_accepted):
        # Time remaining from current point
        time_remaining = keras.ops.abs(stop_time - _time)
        min_step_size = time_remaining / (max_steps - _step)
        max_step_size = time_remaining / keras.ops.maximum(min_steps - _step, 1.0)
        h = keras.ops.sign(_step_size) * keras.ops.clip(keras.ops.abs(_step_size), min_step_size, max_step_size)

        # Take one trial step
        new_state, new_time, new_k1, err = step_fn(
            state=_state,
            time=_time,
            step_size=h,
            k1=_k1,
        )

        # New step size suggestion
        max_abs = None
        for k, v in _state.items():
            m = keras.ops.max(keras.ops.abs(v))
            max_abs = m if max_abs is None else keras.ops.maximum(max_abs, m)
        scale = atol + rtol * max_abs
        error_ratio = err / scale
        new_step_size = h * keras.ops.clip(0.9 * (1.0 / (error_ratio + 1e-12)) ** 0.2, 0.2, 5.0)
        new_step_size = keras.ops.sign(new_step_size) * keras.ops.clip(
            keras.ops.abs(new_step_size), min_step_size, max_step_size
        )

        # Error control
        too_big = keras.ops.greater(error_ratio, 1.0)
        at_min = keras.ops.less_equal(
            keras.ops.abs(h),
            keras.ops.abs(min_step_size),
        )
        accepted = keras.ops.logical_or(keras.ops.logical_not(too_big), at_min)

        updated_state = keras.ops.cond(accepted, lambda: new_state, lambda: _state)
        updated_time = keras.ops.cond(accepted, lambda: new_time, lambda: _time)
        updated_k1 = keras.ops.cond(accepted, lambda: new_k1, lambda: _k1)

        # Step counter: increment only on accepted steps
        updated_step = _step + keras.ops.where(accepted, 1.0, 0.0)
        _count_not_accepted = _count_not_accepted + keras.ops.where(accepted, 0.0, 1.0)

        # For the next iteration, always use the new suggested step size
        return updated_state, updated_time, new_step_size, updated_step, updated_k1, _count_not_accepted

    # Run the adaptive loop
    state, time, step_size, step, k1, count_not_accepted = keras.ops.while_loop(
        cond,
        body,
        [state, start_time, initial_step, step0, k1_0, count_not_accepted],
    )

    if _check_all_nans(state):
        raise RuntimeError(f"All values are NaNs in state during integration at {time}.")

    # Final step to hit stop_time exactly
    time_diff = stop_time - time
    time_remaining = keras.ops.sign(stop_time - start_time) * time_diff
    if keras.ops.all(time_remaining > 0):
        state, time, _, _ = step_fn(
            state=state,
            time=time,
            step_size=time_diff,
            k1=k1,
        )
        step = step + 1.0

    debug(f"Finished integration after {step} steps with {count_not_accepted} rejected steps.")
    return state


def integrate_scipy(
    fn: Callable,
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    scipy_kwargs: dict | None = None,
) -> StateDict:
    import scipy.integrate

    scipy_kwargs = scipy_kwargs or {}
    keys = list(state.keys())
    # convert to tensor before determining the shape in case a number was passed
    shapes = keras.tree.map_structure(lambda x: keras.ops.shape(keras.ops.convert_to_tensor(x)), state)
    adapter = Adapter().concatenate(keys, into="x", axis=-1).convert_dtype(np.float32, np.float64)

    def state_to_vector(state):
        state = keras.tree.map_structure(keras.ops.convert_to_numpy, state)
        # flatten state
        state = keras.tree.map_structure(lambda x: keras.ops.reshape(x, (-1,)), state)
        # apply concatenation
        x = adapter.forward(state)["x"]
        return x

    def vector_to_state(x):
        state = adapter.inverse({"x": x})
        state = {key: keras.ops.reshape(value, shapes[key]) for key, value in state.items()}
        state = keras.tree.map_structure(keras.ops.convert_to_tensor, state)
        return state

    def scipy_wrapper_fn(time, x):
        state = vector_to_state(x)
        time = keras.ops.convert_to_tensor(time, dtype="float32")
        deltas = fn(time, **filter_kwargs(state, fn))
        return state_to_vector(deltas)

    result = scipy.integrate.solve_ivp(
        scipy_wrapper_fn,
        (start_time, stop_time),
        state_to_vector(state),
        **scipy_kwargs,
    )

    result = vector_to_state(result.y[:, -1])
    return result


def integrate(
    fn: Callable,
    state: StateDict,
    start_time: ArrayLike | None = None,
    stop_time: ArrayLike | None = None,
    min_steps: int = 50,
    max_steps: int = 10_000,
    steps: int | Literal["adaptive"] | Tensor | np.ndarray = 100,
    method: str = "rk45",
    **kwargs,
) -> StateDict:
    if isinstance(steps, str) and steps in ["adaptive", "dynamic"]:
        if start_time is None or stop_time is None:
            raise ValueError(
                "Please provide start_time and stop_time for the integration, was "
                f"'start_time={start_time}', 'stop_time={stop_time}'."
            )
        if method == "scipy":
            if min_steps != 10:
                warning("Setting min_steps has no effect for method 'scipy'")
            if max_steps != 10_000:
                warning("Setting max_steps has no effect for method 'scipy'")
            return integrate_scipy(fn, state, start_time, stop_time)
        return integrate_adaptive(fn, state, start_time, stop_time, min_steps, max_steps, method, **kwargs)
    elif isinstance(steps, int):
        if start_time is None or stop_time is None:
            raise ValueError(
                "Please provide start_time and stop_time for the integration, was "
                f"'start_time={start_time}', 'stop_time={stop_time}'."
            )
        return integrate_fixed(fn, state, start_time, stop_time, steps, method, **kwargs)
    elif isinstance(steps, Sequence) or isinstance(steps, np.ndarray) or keras.ops.is_tensor(steps):
        return integrate_scheduled(fn, state, steps, method, **kwargs)
    else:
        raise RuntimeError(f"Type or value of `steps` not understood (steps={steps})")


############ SDE Solvers #############


def generate_noise(z: StateDict, seed: keras.random.SeedGenerator) -> StateDict:
    noise = {
        k: keras.random.normal(keras.ops.shape(val), dtype=keras.ops.dtype(val), seed=seed) for k, val in z.items()
    }
    return noise


def stochastic_adaptive_step_size_controller(
    state,
    drift,
    adaptive_factor: ArrayLike,
    min_step_size: float = -float("inf"),
    max_step_size: float = float("inf"),
) -> ArrayLike:
    """
    Adaptive step size controller based on [1]. Similar to a tamed explicit Euler method when used in Euler-Maruyama.

    Adaptive step sizing uses:
        h = max(1, ||x||**2) / max(1, ||f(x)||**2) * adaptive_factor


    [1] Fang & Giles, Adaptive Euler-Maruyama Method for SDEs with Non-Globally Lipschitz Drift Coefficients (2020)

    Returns
    -------
    New step size.
    """
    state_norms = []
    drift_norms = []
    for key in state.keys():
        state_norms.append(keras.ops.norm(state[key], ord=2, axis=-1))
        drift_norms.append(keras.ops.norm(drift[key], ord=2, axis=-1))
    state_norm = keras.ops.stack(state_norms)
    drift_norm = keras.ops.stack(drift_norms)
    max_state_norm = keras.ops.maximum(
        keras.ops.cast(1.0, dtype=keras.ops.dtype(state_norm)), keras.ops.max(state_norm) ** 2
    )
    max_drift_norm = keras.ops.maximum(
        keras.ops.cast(1.0, dtype=keras.ops.dtype(drift_norm)), keras.ops.max(drift_norm) ** 2
    )
    new_step_size = max_state_norm / max_drift_norm * adaptive_factor
    new_step_size = keras.ops.clip(new_step_size, min_step_size, max_step_size)
    return new_step_size


def euler_maruyama_step(
    drift_fn: Callable,
    diffusion_fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    noise: StateDict,
    use_adaptive_step_size: bool = False,
    min_step_size: float = -float("inf"),
    max_step_size: float = float("inf"),
) -> Union[Tuple[StateDict, ArrayLike, ArrayLike], Tuple[StateDict, ArrayLike, ArrayLike, StateDict]]:
    """
    Performs a single Euler-Maruyama step for stochastic differential equations.

    Args:
        drift_fn: Function computing the drift term f(t, **state).
        diffusion_fn: Function computing the diffusion term g(t, **state).
        state: Current state, mapping variable names to tensors.
        time: Current time scalar tensor.
        step_size: Time increment dt.
        noise: Mapping of variable names to dW noise tensors.
        use_adaptive_step_size: Whether to use adaptive step sizing.
        min_step_size: Minimum allowed step size.
        max_step_size: Maximum allowed step size.

    Returns:
        new_state: Updated state after one Euler-Maruyama step.
        new_time: time + dt.
    """
    # Compute drift and diffusion
    drift = drift_fn(time, **filter_kwargs(state, drift_fn))
    diffusion = diffusion_fn(time, **filter_kwargs(state, diffusion_fn))

    new_step_size = step_size
    if use_adaptive_step_size:
        sign_step = keras.ops.sign(step_size)
        new_step_size = stochastic_adaptive_step_size_controller(
            state=state,
            drift=drift,
            adaptive_factor=max_step_size,
            min_step_size=min_step_size,
            max_step_size=max_step_size,
        )
        new_step_size = sign_step * keras.ops.abs(new_step_size)

    sqrt_step_size = keras.ops.sqrt(keras.ops.abs(new_step_size))

    new_state = {}
    for key, d in drift.items():
        base = state[key] + new_step_size * d
        if key in diffusion:
            base = base + diffusion[key] * sqrt_step_size * noise[key]
        new_state[key] = base

    if use_adaptive_step_size:
        return new_state, time + new_step_size, new_step_size, state
    return new_state, time + new_step_size, new_step_size


def two_step_adaptive_step(
    drift_fn: Callable,
    diffusion_fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    noise: StateDict,
    last_state: StateDict = None,
    use_adaptive_step_size: bool = True,
    min_step_size: float = -float("inf"),
    max_step_size: float = float("inf"),
    e_rel: float = 0.1,
    e_abs: float = None,
    r: float = 0.9,
    adapt_safety: float = 0.9,
) -> Union[
    Tuple[StateDict, ArrayLike, ArrayLike],
    Tuple[StateDict, ArrayLike, ArrayLike, StateDict],
]:
    """
    Performs a single adaptive step for stochastic differential equations based on [1].

    Based on

    This method uses a predictor-corrector approach with error estimation:
    1. Take an Euler-Maruyama step (predictor)
    2. Take another Euler-Maruyama step from the predicted state
    3. Average the two predictions (corrector)
    4. Estimate error and adapt step size

    When step_size reaches min_step_size, steps are always accepted regardless of
    error to ensure progress and termination within max_steps.

    [1] Jolicoeur-Martineau et al. (2021) "Gotta Go Fast When Generating Data with Score-Based Models"

    Args:
        drift_fn: Function computing the drift term f(t, **state).
        diffusion_fn: Function computing the diffusion term g(t, **state).
        state: Current state, mapping variable names to tensors.
        time: Current time scalar tensor.
        step_size: Time increment dt.
        noise: Mapping of variable names to dW noise tensors (pre-scaled by sqrt(dt)).
        last_state: Previous state for error estimation.
        use_adaptive_step_size: Whether to adapt step size.
        min_step_size: Minimum allowed step size.
        max_step_size: Maximum allowed step size.
        e_rel: Relative error tolerance.
        e_abs: Absolute error tolerance. Default assumes standardized targets.
        r: Order of the method for step size adaptation.
        adapt_safety: Safety factor for step size adaptation.

    Returns:
        new_state: Updated state after one adaptive step.
        new_time: time + dt (or time if step rejected).
        new_step_size: Adapted step size for next iteration.
    """
    state_euler, time_mid, _ = euler_maruyama_step(
        drift_fn=drift_fn,
        diffusion_fn=diffusion_fn,
        state=state,
        time=time,
        step_size=step_size,
        min_step_size=min_step_size,
        max_step_size=max_step_size,
        noise=noise,
        use_adaptive_step_size=False,
    )

    # Compute drift and diffusion at new state, but update from old state
    drift_mid = drift_fn(time_mid, **filter_kwargs(state_euler, drift_fn))
    diffusion_mid = diffusion_fn(time_mid, **filter_kwargs(state_euler, diffusion_fn))
    sqrt_step_size = keras.ops.sqrt(keras.ops.abs(step_size))

    state_euler_mid = {}
    for key, d in drift_mid.items():
        base = state[key] + step_size * d
        if key in diffusion_mid:
            base = base + diffusion_mid[key] * sqrt_step_size * noise[key]
        state_euler_mid[key] = base

    # average the two predictions
    state_heun = {}
    for key in state.keys():
        state_heun[key] = 0.5 * (state_euler[key] + state_euler_mid[key])

    # Error estimation
    if use_adaptive_step_size:
        if e_abs is None:
            e_abs = 0.02576  # 1% of 99% CI of standardized unit variance
        # Check if we're at minimum step size - if so, force acceptance
        at_min_step = keras.ops.less_equal(keras.ops.abs(step_size), min_step_size)

        # Compute error tolerance for each component
        e_abs_tensor = keras.ops.cast(e_abs, dtype=keras.ops.dtype(list(state.values())[0]))
        e_rel_tensor = keras.ops.cast(e_rel, dtype=keras.ops.dtype(list(state.values())[0]))

        max_error = keras.ops.cast(0.0, dtype=keras.ops.dtype(list(state.values())[0]))

        for key in state.keys():
            # Local error estimate: difference between Heun and first Euler step
            error_estimate = keras.ops.abs(state_heun[key] - state_euler[key])

            # Tolerance threshold
            delta = keras.ops.maximum(
                e_abs_tensor,
                e_rel_tensor * keras.ops.maximum(keras.ops.abs(state_euler[key]), keras.ops.abs(last_state[key])),
            )

            # Normalized error
            normalized_error = error_estimate / (delta + 1e-10)

            # Maximum error across all components and batch dimensions
            component_max_error = keras.ops.max(normalized_error)
            max_error = keras.ops.maximum(max_error, component_max_error)

        error_scale = 1  # 1/sqrt(n_params)
        E2 = error_scale * max_error

        # Accept step if error is acceptable OR if at minimum step size
        error_acceptable = keras.ops.less_equal(E2, keras.ops.cast(1.0, dtype=keras.ops.dtype(E2)))
        accepted = keras.ops.logical_or(error_acceptable, at_min_step)

        # Adapt step size for next iteration (only if not at minimum)
        # Ensure E2 is not zero to avoid division issues
        E2_safe = keras.ops.maximum(E2, 1e-10)

        # New step size based on error estimate
        adapt_factor = adapt_safety * keras.ops.power(E2_safe, -r)
        new_step_candidate = step_size * adapt_factor

        # Clamp to valid range
        new_step_size = keras.ops.clip(keras.ops.abs(new_step_candidate), min_step_size, max_step_size)
        new_step_size = keras.ops.sign(step_size) * new_step_size

        # Return appropriate state based on acceptance
        new_state = keras.ops.cond(accepted, lambda: state_heun, lambda: state)

        new_time = keras.ops.cond(accepted, lambda: time_mid, lambda: time)

        prev_state = keras.ops.cond(accepted, lambda: state_euler, lambda: state)

        return new_state, new_time, new_step_size, prev_state

    else:
        return state_heun, time_mid, step_size


def compute_levy_area(
    state: StateDict, diffusion: StateDict, noise: StateDict, noise_aux: StateDict, step_size: ArrayLike
) -> StateDict:
    step_size_abs = keras.ops.abs(step_size)
    sqrt_step_size = keras.ops.sqrt(step_size_abs)
    inv_sqrt3 = keras.ops.cast(1.0 / np.sqrt(3.0), dtype=keras.ops.dtype(step_size_abs))

    # Build Lévy area H_k from w_k and Z_k
    H = {}
    for k in state.keys():
        if k in diffusion:
            term1 = 0.5 * step_size_abs * noise[k]
            term2 = 0.5 * step_size_abs * sqrt_step_size * inv_sqrt3 * noise_aux[k]
            H[k] = term1 + term2
        else:
            H[k] = keras.ops.zeros_like(state[k])
    return H


def sea_step(
    drift_fn: Callable,
    diffusion_fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    noise: StateDict,  # standard normals
    noise_aux: StateDict,  # standard normals
) -> Tuple[StateDict, ArrayLike, ArrayLike]:
    """
    Performs a single shifted Euler step for SDEs with additive noise [1].

    Compared to Euler-Maruyama, this evaluates the drift at a shifted state,
    which improves the local error and the global error constant for additive noise.

    The scheme is
        X_{n+1} = X_n + f(t_n, X_n + g(t_n) * (0.5 * ΔW_n + ΔH_n) * h + g(t_n) * ΔW_n

    [1] Foster et al., "High order splitting methods for SDEs satisfying a commutativity condition" (2023)
    Args:
        drift_fn: Function computing the drift term f(t, **state).
        diffusion_fn: Function computing the diffusion term g(t, **state).
        state: Current state, mapping variable names to tensors.
        time: Current time scalar tensor.
        step_size: Time increment dt.
        noise: Mapping of variable names to dW noise tensors.
        noise_aux: Mapping of variable names to auxiliary noise.

    Returns:
        new_state: Updated state after one SEA step.
        new_time: time + dt.
    """
    # Compute diffusion
    diffusion = diffusion_fn(time, **filter_kwargs(state, diffusion_fn))
    sqrt_step_size = keras.ops.sqrt(keras.ops.abs(step_size))

    la = compute_levy_area(state=state, diffusion=diffusion, noise=noise, noise_aux=noise_aux, step_size=step_size)

    # Build shifted state: X_shift = X + g * (0.5 * ΔW + ΔH)
    shifted_state = {}
    for key, x in state.items():
        if key in diffusion:
            shifted_state[key] = x + diffusion[key] * (0.5 * sqrt_step_size * noise[key] + la[key])
        else:
            shifted_state[key] = x

    # Drift evaluated at shifted state
    drift_shifted = drift_fn(time, **filter_kwargs(shifted_state, drift_fn))

    # Final update
    new_state = {}
    for key, d in drift_shifted.items():
        base = state[key] + step_size * d
        if key in diffusion:
            base = base + diffusion[key] * sqrt_step_size * noise[key]
        new_state[key] = base

    return new_state, time + step_size, step_size


def shark_step(
    drift_fn: Callable,
    diffusion_fn: Callable,
    state: StateDict,
    time: ArrayLike,
    step_size: ArrayLike,
    noise: StateDict,
    noise_aux: StateDict,
) -> Tuple[StateDict, ArrayLike, ArrayLike]:
    """
    Shifted Additive noise Runge Kutta (SHARK) for additive SDEs [1]. Makes two evaluations of the drift and diffusion
    per step and has a strong order 1.5.

    SHARK method as specified:

        1)  ỹ_k = y_k + g(y_k) H_k
        2)  ỹ_{k+5/6} = ỹ_k + (5/6)[ f(ỹ_k) h + g(ỹ_k) W_k ]
        3)  y_{k+1} = y_k
                     + (2/5) f(ỹ_k) h
                     + (3/5) f(ỹ_{k+5/6}) h
                     + g(ỹ_k) ( 2/5 W_k +  6/5 H_k )
                     + g(ỹ_{k+5/6}) ( 3/5 W_k -  6/5 H_k )

        with
            H_k = 0.5 * |h| * W_k + (|h| ** 1.5) / (2 * sqrt(3)) * Z_k

    [1] Foster et al., "High order splitting methods for SDEs satisfying a commutativity condition" (2023)

    Args:
        drift_fn: Function computing the drift term f(t, **state).
        diffusion_fn: Function computing the diffusion term g(t, **state).
        state: Current state, mapping variable names to tensors.
        time: Current time scalar tensor.
        step_size: Time increment dt.
        noise: Mapping of variable names to dW noise tensors.
        noise_aux: Mapping of variable names to auxiliary noise.

    Returns:
        new_state: Updated state after one SHARK step.
        new_time: time + dt.
    """
    h = step_size
    t = time
    h_mag = keras.ops.abs(h)
    sqrt_h_mag = keras.ops.sqrt(h_mag)

    diffusion = diffusion_fn(t, **filter_kwargs(state, diffusion_fn))

    la = compute_levy_area(state=state, diffusion=diffusion, noise=noise, noise_aux=noise_aux, step_size=step_size)

    # === 1) shifted initial state ===
    y_tilde_k = {}
    for k in state.keys():
        if k in diffusion:
            y_tilde_k[k] = state[k] + diffusion[k] * la[k]
        else:
            y_tilde_k[k] = state[k]

    # === evaluate drift and diffusion at ỹ_k ===
    f_tilde_k = drift_fn(t, **filter_kwargs(y_tilde_k, drift_fn))
    g_tilde_k = diffusion_fn(t, **filter_kwargs(y_tilde_k, diffusion_fn))

    # === 2) internal stage at 5/6 ===
    y_tilde_mid = {}
    for k in state.keys():
        drift_part = (5.0 / 6.0) * f_tilde_k[k] * h
        if k in g_tilde_k:
            sto_part = (5.0 / 6.0) * g_tilde_k[k] * sqrt_h_mag * noise[k]
        else:
            sto_part = keras.ops.zeros_like(state[k])
        y_tilde_mid[k] = y_tilde_k[k] + drift_part + sto_part

    # === evaluate drift and diffusion at ỹ_(k+5/6) ===
    f_tilde_mid = drift_fn(t + 5.0 / 6.0 * h, **filter_kwargs(y_tilde_mid, drift_fn))
    g_tilde_mid = diffusion_fn(t + 5.0 / 6.0 * h, **filter_kwargs(y_tilde_mid, diffusion_fn))

    # === 3) final update ===
    new_state = {}
    for k in state.keys():
        # deterministic weights
        det = state[k] + (2.0 / 5.0) * f_tilde_k[k] * h + (3.0 / 5.0) * f_tilde_mid[k] * h

        # stochastic parts
        sto1 = (
            g_tilde_k[k] * ((2.0 / 5.0) * sqrt_h_mag * noise[k] + (6.0 / 5.0) * la[k])
            if k in g_tilde_k
            else keras.ops.zeros_like(det)
        )
        sto2 = (
            g_tilde_mid[k] * ((3.0 / 5.0) * sqrt_h_mag * noise[k] - (6.0 / 5.0) * la[k])
            if k in g_tilde_mid
            else keras.ops.zeros_like(det)
        )

        new_state[k] = det + sto1 + sto2

    return new_state, t + h, h


def _apply_corrector(
    new_state: StateDict,
    new_time: ArrayLike,
    i: ArrayLike,
    corrector_steps: int,
    score_fn: Optional[Callable],
    corrector_noise_history: StateDict | None,
    seed: keras.random.SeedGenerator,
    step_size_factor: ArrayLike = 0.01,
    noise_schedule=None,
) -> StateDict:
    """Helper function to apply corrector steps [1].

    [1] Song et al., "Score-Based Generative Modeling through Stochastic Differential Equations" (2020)
    """
    for j in range(corrector_steps):
        score = score_fn(new_time, **filter_kwargs(new_state, score_fn))
        if corrector_noise_history is None:
            _z_corr = generate_noise(new_state, seed=seed)
        else:
            _z_corr = {k: val[i, j] for k, val in corrector_noise_history.items()}

        log_snr_t = noise_schedule.get_log_snr(t=new_time, training=False)
        alpha_t, _ = noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)

        for k in new_state.keys():
            if k in score:
                # Calculate required norms for Langevin step
                z_norm = keras.ops.norm(_z_corr[k], axis=-1, keepdims=True)
                score_norm = keras.ops.norm(score[k], axis=-1, keepdims=True)
                score_norm = keras.ops.maximum(score_norm, 1e-8)

                # Compute step size for the Langevin update
                e = 2.0 * alpha_t * (step_size_factor * z_norm / score_norm) ** 2

                # Annealed Langevin Dynamics update
                new_state[k] = new_state[k] + e * score[k] + keras.ops.sqrt(2.0 * e) * _z_corr[k]
    return new_state


def integrate_stochastic_fixed(
    step_fn: Callable,
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    steps: int,
    min_step_size: ArrayLike,
    max_step_size: ArrayLike,
    z_history: StateDict | None,
    z_extra_history: StateDict | None,
    score_fn: Optional[Callable],
    step_size_factor: ArrayLike,
    corrector_noise_history: StateDict | None,
    seed: keras.random.SeedGenerator,
    corrector_steps: int = 0,
    noise_schedule=None,
) -> StateDict:
    """
    Performs fixed-step SDE integration.
    """
    initial_step = (stop_time - start_time) / float(steps)

    def cond(_loop_var, _loop_state, _loop_time, _loop_step):
        all_nans = _check_all_nans(_loop_state)
        end_now = keras.ops.less(_loop_var, steps)
        return keras.ops.logical_and(~all_nans, end_now)

    def body(_i, _current_state, _current_time, _current_step):
        # Determine step size: either the constant size or the remainder to reach stop_time
        remaining = keras.ops.abs(stop_time - _current_time)
        sign = keras.ops.sign(_current_step)
        dt_mag = keras.ops.minimum(keras.ops.abs(_current_step), remaining)
        dt = sign * dt_mag

        # Generate noise increment
        if z_history is None:
            _noise_i = generate_noise(_current_state, seed=seed)
        else:
            _noise_i = {k: val[_i] for k, val in z_history.items()}
        _noise_extra_i = None
        if z_extra_history is not None:
            if len(z_extra_history) == 0:
                _noise_extra_i = generate_noise(_current_state, seed=seed)
            else:
                _noise_extra_i = {k: val[_i] for k, val in z_history.items()}

        step_fn_additional_args = dict(
            noise_aux=_noise_extra_i,
            min_step_size=min_step_size,
            max_step_size=keras.ops.minimum(max_step_size, remaining),
            use_adaptive_step_size=False,
        )
        new_state, new_time, new_step = step_fn(
            state=_current_state,
            time=_current_time,
            step_size=dt,
            noise=_noise_i,
            **filter_kwargs(step_fn_additional_args, step_fn),
        )

        if corrector_steps > 0:
            new_state = _apply_corrector(
                new_state=new_state,
                new_time=new_time,
                i=_i,
                corrector_steps=corrector_steps,
                score_fn=score_fn,
                noise_schedule=noise_schedule,
                step_size_factor=step_size_factor,
                corrector_noise_history=corrector_noise_history,
                seed=seed,
            )
        return _i + 1, new_state, new_time, initial_step

    _, final_state, final_time, _ = keras.ops.while_loop(
        cond,
        body,
        [0, state, start_time, initial_step],
    )
    if _check_all_nans(final_state):
        raise RuntimeError(f"All values are NaNs in state during integration at {final_time}.")

    return final_state


def integrate_stochastic_adaptive(
    step_fn: Callable,
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    max_steps: int,
    min_step_size: ArrayLike,
    max_step_size: ArrayLike,
    initial_step: ArrayLike,
    z_history: StateDict | None,
    z_extra_history: StateDict | None,
    score_fn: Optional[Callable],
    step_size_factor: ArrayLike,
    seed: keras.random.SeedGenerator,
    corrector_noise_history: StateDict | None,
    corrector_steps: int = 0,
    noise_schedule=None,
) -> StateDict:
    """
    Performs adaptive-step SDE integration.
    """
    initial_loop_state = (keras.ops.zeros((), dtype="int32"), state, start_time, initial_step, state)
    if keras.backend.backend() == "jax":
        seed = None  # not needed, noise is generated upfront
    else:
        seed_body = seed

    def cond(i, current_state, current_time, current_step, last_state):
        time_remaining = keras.ops.sign(stop_time - start_time) * (stop_time - (current_time + current_step))
        all_nans = _check_all_nans(current_state)
        end_now = keras.ops.logical_and(keras.ops.all(time_remaining > 0), keras.ops.less(i, max_steps))
        return keras.ops.logical_and(~all_nans, end_now)

    def body_adaptive(_i, _current_state, _current_time, _current_step, _last_state):
        # Step Size Control
        remaining = keras.ops.abs(stop_time - _current_time)
        sign = keras.ops.sign(_current_step)
        # Ensure the next step does not overshoot the stop_time
        dt_mag = keras.ops.minimum(keras.ops.abs(_current_step), remaining)
        dt = sign * dt_mag

        if z_history is None:
            _noise_i = generate_noise(_current_state, seed=seed_body)
        else:
            _noise_i = {k: val[_i] for k, val in z_history.items()}

        _noise_extra_i = None
        if z_extra_history is not None:
            if len(z_extra_history) == 0:
                _noise_extra_i = generate_noise(_current_state, seed=seed_body)
            else:
                _noise_extra_i = {k: val[_i] for k, val in z_history.items()}

        step_fn_additional_args = dict(
            last_state=_last_state,
            noise_aux=_noise_extra_i,
        )
        new_state, new_time, new_step, _new_current_state = step_fn(
            state=_current_state,
            time=_current_time,
            step_size=dt,
            min_step_size=min_step_size,
            max_step_size=keras.ops.minimum(max_step_size, remaining),
            noise=_noise_i,
            use_adaptive_step_size=True,
            **filter_kwargs(step_fn_additional_args, step_fn),
        )

        if corrector_steps > 0:
            new_state = _apply_corrector(
                new_state=new_state,
                new_time=new_time,
                i=_i,
                corrector_steps=corrector_steps,
                score_fn=score_fn,
                noise_schedule=noise_schedule,
                step_size_factor=step_size_factor,
                corrector_noise_history=corrector_noise_history,
                seed=seed_body,
            )

        return _i + 1, new_state, new_time, new_step, _new_current_state

    # Execute the adaptive loop
    final_counter, final_state, final_time, _, final_k1 = keras.ops.while_loop(cond, body_adaptive, initial_loop_state)

    if _check_all_nans(final_state):
        raise RuntimeError(f"All values are NaNs in state during integration at {final_time}.")

    # Final step to hit stop_time exactly
    time_diff = stop_time - final_time
    time_remaining = keras.ops.sign(stop_time - start_time) * time_diff
    if keras.ops.all(time_remaining > 0):
        noise_final = generate_noise(final_state, seed=seed)
        noise_extra_final = None
        if z_extra_history is not None and len(z_extra_history) > 0:
            noise_extra_final = generate_noise(final_state, seed=seed)

        step_fn_additional_args_final = dict(
            noise_aux=noise_extra_final,
            last_state=final_k1,
        )
        final_state, _, _ = step_fn(
            state=final_state,
            time=final_time,
            step_size=time_diff,
            min_step_size=min_step_size,
            max_step_size=time_remaining,
            noise=noise_final,
            use_adaptive_step_size=False,
            **filter_kwargs(step_fn_additional_args_final, step_fn),
        )
        final_counter = final_counter + 1

    debug(f"Finished integration after {final_counter}.")
    return final_state


def integrate_langevin(
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    steps: int,
    z_history: StateDict | None,
    score_fn: Callable,
    noise_schedule,
    seed: keras.random.SeedGenerator,
    corrector_noise_history: StateDict | None,
    step_size_factor: ArrayLike = 0.01,
    corrector_steps: int = 0,
) -> StateDict:
    """
    Annealed Langevin dynamics using the given score_fn and noise_schedule [1].

    At each step i with time t_i, performs for every state component k:
        state_k <- state_k + e * score_k + sqrt(2 * e) * z

    Times are stepped linearly from start_time to stop_time.

    [1]  Song et al., "Generative Modeling by Estimating Gradients of the Data Distribution" (2020)
    """

    if steps <= 0:
        raise ValueError("Number of Langevin steps must be positive.")
    if score_fn is None or noise_schedule is None:
        raise ValueError("score_fn and noise_schedule must be provided.")

    # Linear time grid
    dt = (stop_time - start_time) / float(steps)
    effective_factor = step_size_factor * 100 / np.sqrt(steps)

    def cond(_loop_var, _loop_state, _loop_time):
        all_nans = _check_all_nans(_loop_state)
        end_now = keras.ops.less(_loop_var, steps)
        return keras.ops.logical_and(~all_nans, end_now)

    def body(_i, _loop_state, _loop_time):
        # score at current time
        score = score_fn(_loop_time, **filter_kwargs(_loop_state, score_fn))

        # noise schedule
        log_snr_t = noise_schedule.get_log_snr(t=_loop_time, training=False)
        _, sigma_t = noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)

        new_state: StateDict = {}
        if z_history is None:
            z_history_i = generate_noise(_loop_state, seed=seed)
        else:
            z_history_i = {k: val[_i] for k, val in z_history.items()}
        for k in _loop_state.keys():
            s_k = score.get(k, None)
            if s_k is None:
                new_state[k] = _loop_state[k]
                continue

            e = effective_factor * sigma_t**2
            new_state[k] = _loop_state[k] + e * s_k + keras.ops.sqrt(2.0 * e) * z_history_i[k]

        new_time = _loop_time + dt

        if corrector_steps > 0:
            new_state = _apply_corrector(
                new_state=new_state,
                new_time=new_time,
                i=_i,
                corrector_steps=corrector_steps,
                score_fn=score_fn,
                noise_schedule=noise_schedule,
                step_size_factor=step_size_factor,
                corrector_noise_history=corrector_noise_history,
                seed=seed,
            )

        return _i + 1, new_state, new_time

    _, final_state, final_time = keras.ops.while_loop(
        cond,
        body,
        (0, state, start_time),
    )
    if _check_all_nans(final_state):
        raise RuntimeError(f"All values are NaNs in state during integration at {final_time}.")
    return final_state


def integrate_stochastic(
    drift_fn: Callable,
    diffusion_fn: Callable,
    state: StateDict,
    start_time: ArrayLike,
    stop_time: ArrayLike,
    seed: keras.random.SeedGenerator,
    steps: int | Literal["adaptive"] = 100,
    method: str = "euler_maruyama",
    min_steps: int = 50,
    max_steps: int = 1_000,
    score_fn: Callable = None,
    corrector_steps: int = 0,
    noise_schedule=None,
    step_size_factor: ArrayLike = 0.01,
    **kwargs,
) -> StateDict:
    """
    Integrate a stochastic differential equation from ``start_time`` to ``stop_time``.

    This function dispatches to either fixed-step or adaptive-step integration logic,
    depending on the selected integration ``method`` and the value of ``steps``.

    Parameters
    ----------
    drift_fn : callable
        Function computing the drift term of the SDE. It should accept the current
        state and time as inputs.
    diffusion_fn : callable
        Function computing the diffusion term of the SDE. It should accept the current
        state and time as inputs.
    state : StateDict
        Dictionary containing the initial state of the system.
    start_time : array-like
        Starting time for integration.
    stop_time : array-like
        Ending time for integration.
    seed : keras.random.SeedGenerator
        Random seed generator used for noise generation.
    steps : int or {'adaptive'}, optional
        Number of integration steps for fixed-step integration, or ``'adaptive'``
        to enable adaptive step sizing. Adaptive integration is only supported by
        the ``'shark'`` method. Default is ``100``.
    method : str, optional
        Integration method to use (e.g., ``'euler_maruyama'`` or ``'shark'``).
        Default is ``'euler_maruyama'``.
    min_steps : int, optional
        Minimum number of steps for adaptive integration. Default is ``50``.
    max_steps : int, optional
        Maximum number of steps for adaptive integration. Noise is pre-generated
        up to this number of steps, which may increase memory usage. Default is
        ``1000``.
    score_fn : callable, optional
        Score function used for predictor–corrector sampling. If ``None``,
        no corrector step is applied.
    corrector_steps : int, optional
        Number of corrector steps applied after each predictor step. Default is ``0``.
    noise_schedule : object, optional
        Noise schedule object used to compute ``alpha_t`` during the corrector step.
        Required if ``corrector_steps > 0``.
    step_size_factor : array-like, optional
        Scaling factor applied to the corrector step size. Default is ``0.01``.
    **kwargs
        Additional keyword arguments passed to the underlying step function.

    Returns
    -------
    StateDict
        Final state dictionary after integration.
    """

    is_adaptive = isinstance(steps, str) and steps in ["adaptive", "dynamic"]

    if is_adaptive:
        if start_time is None or stop_time is None:
            raise ValueError("Please provide start_time and stop_time for adaptive integration.")
        if min_steps <= 0 or max_steps <= 0 or max_steps < min_steps:
            raise ValueError("min_steps and max_steps must be positive, and max_steps >= min_steps.")

        loop_steps = max_steps
        initial_step = (stop_time - start_time) / float(min_steps)
        span_mag = keras.ops.abs(stop_time - start_time)
        min_step_size = span_mag / keras.ops.cast(max_steps, span_mag.dtype)
        max_step_size = span_mag / keras.ops.cast(min_steps, span_mag.dtype)
    else:
        if steps <= 0:
            raise ValueError("Number of steps must be positive.")
        loop_steps = int(steps)
        initial_step = (stop_time - start_time) / float(loop_steps)
        # For fixed step, min/max step size are just the fixed step size
        min_step_size, max_step_size = initial_step, initial_step

    # Pre-generate corrector noise if requested
    corrector_noise_history = None
    if corrector_steps > 0:
        if score_fn is None or noise_schedule is None:
            raise ValueError("Please provide both score_fn and noise_schedule when using corrector_steps > 0.")
        if keras.backend.backend() == "jax":
            corrector_noise_history = {}
            for key, val in state.items():
                shape = keras.ops.shape(val)
                corrector_noise_history[key] = keras.random.normal(
                    (loop_steps, corrector_steps, *shape), dtype=keras.ops.dtype(val), seed=seed
                )

    match method:
        case "euler_maruyama":
            step_fn_raw = euler_maruyama_step
        case "sea":
            step_fn_raw = sea_step
            if is_adaptive:
                raise ValueError("SEA SDE solver does not support adaptive steps.")
        case "shark":
            step_fn_raw = shark_step
            if is_adaptive:
                raise ValueError("SHARK SDE solver does not support adaptive steps.")
        case "two_step_adaptive":
            step_fn_raw = two_step_adaptive_step
        case "langevin":
            if is_adaptive:
                raise ValueError("Langevin sampling does not support adaptive steps.")

            z_history = None
            if keras.backend.backend() == "jax":
                warning(f"JAX backend needs to preallocate random samples for 'max_steps={max_steps}'.")
                z_history = {}
                for key, val in state.items():
                    shape = keras.ops.shape(val)
                    z_history[key] = keras.random.normal((loop_steps, *shape), dtype=keras.ops.dtype(val), seed=seed)

            return integrate_langevin(
                state=state,
                start_time=start_time,
                stop_time=stop_time,
                steps=loop_steps,
                z_history=z_history,
                score_fn=score_fn,
                noise_schedule=noise_schedule,
                step_size_factor=step_size_factor,
                corrector_steps=corrector_steps,
                corrector_noise_history=corrector_noise_history,
                seed=seed,
            )
        case other:
            raise TypeError(f"Invalid integration method: {other!r}")

    # Partial the step function with common arguments
    step_fn = partial(step_fn_raw, drift_fn=drift_fn, diffusion_fn=diffusion_fn, **filter_kwargs(kwargs, step_fn_raw))

    # Pre-generate standard normals for the predictor step (up to max_steps)
    z_history = None
    z_extra_history = None if method not in ["sea", "shark"] else {}
    if keras.backend.backend() == "jax":
        warning(f"JAX backend needs to preallocate random samples for 'max_steps={max_steps}'.")
        z_history = {}
        for key, val in state.items():
            shape = keras.ops.shape(val)
            z_history[key] = keras.random.normal((loop_steps, *shape), dtype=keras.ops.dtype(val), seed=seed)
            if method in ["sea", "shark"]:
                z_extra_history[key] = keras.random.normal((loop_steps, *shape), dtype=keras.ops.dtype(val), seed=seed)

    if is_adaptive:
        return integrate_stochastic_adaptive(
            step_fn=step_fn,
            state=state,
            start_time=start_time,
            stop_time=stop_time,
            max_steps=max_steps,
            min_step_size=min_step_size,
            max_step_size=max_step_size,
            initial_step=initial_step,
            z_history=z_history,
            z_extra_history=z_extra_history,
            corrector_steps=corrector_steps,
            score_fn=score_fn,
            noise_schedule=noise_schedule,
            step_size_factor=step_size_factor,
            corrector_noise_history=corrector_noise_history,
            seed=seed,
        )
    else:
        return integrate_stochastic_fixed(
            step_fn=step_fn,
            state=state,
            start_time=start_time,
            stop_time=stop_time,
            min_step_size=min_step_size,
            max_step_size=max_step_size,
            steps=loop_steps,
            z_history=z_history,
            z_extra_history=z_extra_history,
            corrector_steps=corrector_steps,
            score_fn=score_fn,
            noise_schedule=noise_schedule,
            step_size_factor=step_size_factor,
            corrector_noise_history=corrector_noise_history,
            seed=seed,
        )
