"""Flow-matching ODE/SDE sampler for ACE-Step diffusion in MLX.

Implements the sampling primitives for flow-matching diffusion models:
timestep scheduling, Euler ODE steps, SDE steps with re-noising,
classifier-free guidance, and the main denoising loop.

All operations use mlx.core (mx) arrays. No model weights needed --
this is pure math.
"""

from typing import Callable, Optional, Tuple

import mlx.core as mx


# Pre-defined timestep schedules for turbo mode (fix_nfe=8).
# Keys are shift values; values are the 8 timesteps (excluding final 0).
# These come directly from the reference ACE-Step v1.5 implementation.
_TURBO_SCHEDULES = {
    1.0: [1.0, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125],
    2.0: [
        1.0,
        0.9333333333333333,
        0.8571428571428571,
        0.7692307692307693,
        0.6666666666666666,
        0.5454545454545454,
        0.4,
        0.2222222222222222,
    ],
    3.0: [
        1.0,
        0.9545454545454546,
        0.9,
        0.8333333333333334,
        0.75,
        0.6428571428571429,
        0.5,
        0.3,
    ],
}

# All valid discrete timestep values (union of the three schedules above).
VALID_TIMESTEPS = sorted(
    set(t for schedule in _TURBO_SCHEDULES.values() for t in schedule),
    reverse=True,
)

VALID_SHIFTS = [1.0, 2.0, 3.0]


def time_schedule(num_steps: int, shift: float = 3.0) -> mx.array:
    """Compute the timestep schedule for flow-matching sampling.

    For the turbo model with known shift values (1, 2, 3) and 8 steps,
    returns the pre-defined schedule. For arbitrary parameters, computes
    the shifted schedule analytically:

        t_shifted = shift * t / (1 + (shift - 1) * t)

    where t is a linear ramp from 1 down to 0 (excluding the final 0).

    Args:
        num_steps: Number of denoising steps (NFE).
        shift: Time-shift parameter. Higher values front-load the schedule,
            spending more time at higher noise levels.

    Returns:
        mx.array of shape [num_steps] with timesteps from ~1.0 down to >0.
    """
    # Check for pre-defined turbo schedules (exact match on shift and steps)
    if num_steps == 8 and shift in _TURBO_SCHEDULES:
        return mx.array(_TURBO_SCHEDULES[shift])

    # General formula: linspace from 1 to 0, drop the trailing 0, apply shift
    t = mx.linspace(1.0, 0.0, num_steps + 1)[:-1]

    if shift == 1.0:
        return t

    t_shifted = shift * t / (1.0 + (shift - 1.0) * t)
    return t_shifted


def prepare_noise(
    shape: Tuple[int, ...],
    seed: Optional[int] = None,
    dtype: mx.Dtype = mx.float32,
) -> mx.array:
    """Generate initial Gaussian noise for the diffusion process.

    Args:
        shape: Shape of the noise tensor, e.g. (batch, frames, channels).
        seed: Optional random seed for reproducibility. If None, uses
            the current MLX random state.
        dtype: Output dtype (default float32).

    Returns:
        mx.array of Gaussian noise with the given shape and dtype.
    """
    if seed is not None:
        mx.random.seed(seed)
    return mx.random.normal(shape).astype(dtype)


def euler_step(
    xt: mx.array,
    vt: mx.array,
    t_curr: float,
    t_next: float,
) -> mx.array:
    """Single Euler ODE step for flow matching.

    The flow-matching ODE is dx/dt = -v, integrated from t_curr to t_next:

        x_{next} = x_t - v_t * (t_curr - t_next)

    Args:
        xt: Current noisy sample, shape [B, T, C].
        vt: Predicted velocity, same shape as xt.
        t_curr: Current timestep (higher = noisier).
        t_next: Next timestep (lower = cleaner).

    Returns:
        Updated sample at t_next.
    """
    dt = t_curr - t_next
    return xt - vt * dt


def get_x0_from_noise(
    zt: mx.array,
    vt: mx.array,
    t: mx.array,
) -> mx.array:
    """Recover the clean signal estimate from noisy sample and velocity.

    In flow matching, x_t = t * noise + (1 - t) * x_0, and the model
    predicts v = noise - x_0. Therefore:

        x_0 = x_t - v_t * t

    Args:
        zt: Noisy sample at timestep t, shape [B, T, C].
        vt: Predicted velocity, same shape as zt.
        t: Timestep value(s). Can be a scalar mx.array or shape [B].
            Will be broadcast to match zt dimensions.

    Returns:
        Estimated clean signal x_0.
    """
    # Broadcast t to [B, 1, 1] if it's a 1-D batch vector
    if t.ndim == 1:
        t = mx.expand_dims(t, axis=(-1, -2))
    elif t.ndim == 0:
        pass  # scalar broadcasts naturally
    return zt - vt * t


def sde_step(
    xt: mx.array,
    vt: mx.array,
    t_curr: float,
    t_next: float,
) -> mx.array:
    """Single SDE step: predict clean x0 then re-noise at t_next.

    This is the stochastic variant of the sampler. Instead of
    deterministically stepping along the ODE, it:
    1. Estimates x_0 from the current (xt, vt, t_curr).
    2. Re-noises x_0 to the next timestep: x_{next} = t_next * eps + (1 - t_next) * x_0.

    This injects fresh randomness at each step, which can improve
    diversity at the cost of some coherence.

    Args:
        xt: Current noisy sample, shape [B, T, C].
        vt: Predicted velocity, same shape as xt.
        t_curr: Current timestep.
        t_next: Next timestep.

    Returns:
        Updated sample at t_next (stochastic).
    """
    # Build a broadcastable timestep tensor matching xt batch dim
    bsz = xt.shape[0]
    t_curr_arr = mx.full((bsz,), t_curr)

    x0 = get_x0_from_noise(xt, vt, t_curr_arr)

    # Re-noise: x_next = t_next * noise + (1 - t_next) * x0
    noise = mx.random.normal(xt.shape).astype(xt.dtype)
    xt_next = t_next * noise + (1.0 - t_next) * x0
    return xt_next


def apply_cfg(
    pred_cond: mx.array,
    pred_uncond: mx.array,
    scale: float,
) -> mx.array:
    """Apply classifier-free guidance.

    Blends conditional and unconditional predictions:

        output = pred_uncond + scale * (pred_cond - pred_uncond)

    With scale=1.0 this returns pred_cond exactly (no guidance).
    The turbo model was distilled with CFG baked in, so this is
    primarily useful for the base (non-turbo) model.

    Args:
        pred_cond: Conditional model prediction.
        pred_uncond: Unconditional model prediction.
        scale: Guidance scale (1.0 = no guidance, >1.0 = stronger).

    Returns:
        Guided prediction.
    """
    return pred_uncond + scale * (pred_cond - pred_uncond)


def sample_loop(
    denoise_fn: Callable[[mx.array, float], mx.array],
    noise: mx.array,
    schedule: mx.array,
    method: str = "ode",
    cfg_scale: Optional[float] = None,
    denoise_fn_uncond: Optional[Callable[[mx.array, float], mx.array]] = None,
    callback: Optional[Callable[[int, int, mx.array], None]] = None,
) -> mx.array:
    """Main denoising loop over the timestep schedule.

    Iterates through the schedule from t~1.0 (noise) to t~0 (clean),
    applying either Euler ODE steps or SDE steps. On the final step,
    directly recovers x_0 instead of stepping.

    Args:
        denoise_fn: Callable(xt, t) -> vt that runs the diffusion model.
            Takes the current noisy sample and scalar timestep, returns
            the predicted velocity. If cfg_scale is used, this should be
            the *conditional* prediction function.
        noise: Initial noise tensor, shape [B, T, C].
        schedule: Timestep schedule from time_schedule(), shape [num_steps].
        method: Sampling method, either "ode" (deterministic Euler) or
            "sde" (stochastic re-noising).
        cfg_scale: If set, applies classifier-free guidance with this scale.
            Requires denoise_fn_uncond to also be provided.
        denoise_fn_uncond: Unconditional prediction function for CFG.
            Only used when cfg_scale is not None.
        callback: Optional callback(step_idx, total_steps, xt) called
            after each step for progress tracking.

    Returns:
        Denoised sample x_0, same shape as noise.

    Raises:
        ValueError: If method is not "ode" or "sde".
        ValueError: If cfg_scale is set but denoise_fn_uncond is missing.
    """
    if method not in ("ode", "sde"):
        raise ValueError(f"Unknown sampling method: {method!r}. Use 'ode' or 'sde'.")

    if cfg_scale is not None and denoise_fn_uncond is None:
        raise ValueError("cfg_scale requires denoise_fn_uncond to be provided.")

    num_steps = schedule.shape[0]
    xt = noise

    for step_idx in range(num_steps):
        t_curr = schedule[step_idx].item()

        # Get velocity prediction from the model
        vt = denoise_fn(xt, t_curr)

        # Apply classifier-free guidance if requested
        if cfg_scale is not None:
            vt_uncond = denoise_fn_uncond(xt, t_curr)
            vt = apply_cfg(vt, vt_uncond, cfg_scale)

        # Final step: directly compute x0
        if step_idx == num_steps - 1:
            bsz = xt.shape[0]
            t_arr = mx.full((bsz,), t_curr)
            xt = get_x0_from_noise(xt, vt, t_arr)
            mx.eval(xt)
            if callback is not None:
                callback(step_idx, num_steps, xt)
            break

        # Intermediate step
        t_next = schedule[step_idx + 1].item()

        if method == "sde":
            xt = sde_step(xt, vt, t_curr, t_next)
        else:
            xt = euler_step(xt, vt, t_curr, t_next)

        # Force computation between steps (MLX is lazy-evaluated)
        mx.eval(xt)

        if callback is not None:
            callback(step_idx, num_steps, xt)

    return xt
