"""Noise schedule functions for MDLM diffusion.

All functions operate on PyTorch tensors and are pure (no global state).

Convention: alpha(t) is the fraction of tokens that remain *unmasked*.
  - alpha(0) = 1.0  (fully clean)
  - alpha(1) = 0.0  (fully masked)
"""

from __future__ import annotations

import math
from collections.abc import Callable

import torch
from torch import Tensor


def linear_schedule(t: Tensor) -> Tensor:
    """Linear noise schedule: alpha(t) = 1 - t.

    Args:
        t: Diffusion time in [0, 1]. Any shape.

    Returns:
        Retention probability alpha_t, same shape as *t*.
    """
    return 1.0 - t


def cosine_schedule(t: Tensor) -> Tensor:
    """Cosine noise schedule: alpha(t) = cos(pi/2 * t).

    MDLM Appendix E.1 eq (92) ("Cosine"); the same function the craftax
    repo names "cosine".

    Args:
        t: Diffusion time in [0, 1]. Any shape.

    Returns:
        Retention probability alpha_t, same shape as *t*.
    """
    return torch.cos(t * (math.pi / 2.0))


def cosine_sq_schedule(t: Tensor) -> Tensor:
    """Cosine-squared noise schedule: alpha(t) = cos(pi/2 * t)^2.

    MDLM Appendix E.1 eq (91) ("Cosine Squared", after Nichol & Dhariwal).
    Previously registered under the name "cosine" in this repo; renamed so
    the label "cosine" denotes the same function in both repos.

    Args:
        t: Diffusion time in [0, 1]. Any shape.

    Returns:
        Retention probability alpha_t, same shape as *t*.
    """
    return torch.cos(t * (math.pi / 2.0)) ** 2


def linear_schedule_deriv(t: Tensor) -> Tensor:
    """Analytic d(alpha)/dt for the linear schedule."""
    return torch.full_like(t, -1.0)


def cosine_schedule_deriv(t: Tensor) -> Tensor:
    """Analytic d(alpha)/dt for the cosine schedule."""
    return -(math.pi / 2.0) * torch.sin(t * (math.pi / 2.0))


def cosine_sq_schedule_deriv(t: Tensor) -> Tensor:
    """Analytic d(alpha)/dt for the cosine-squared schedule."""
    return -(math.pi / 2.0) * torch.sin(t * math.pi)


_SCHEDULE_MAP: dict[str, Callable[[Tensor], Tensor]] = {
    "linear": linear_schedule,
    "cosine": cosine_schedule,
    "cosine_sq": cosine_sq_schedule,
}

_DERIV_BY_FN: dict[Callable[[Tensor], Tensor], Callable[[Tensor], Tensor]] = {
    linear_schedule: linear_schedule_deriv,
    cosine_schedule: cosine_schedule_deriv,
    cosine_sq_schedule: cosine_sq_schedule_deriv,
}


def get_schedule_deriv_for(
    schedule_fn: Callable[[Tensor], Tensor],
) -> Callable[[Tensor], Tensor]:
    """Analytic derivative for a registered schedule function.

    The NELBO weight uses the analytic d(alpha)/dt as stated in
    MDLM eq (10) / Shi eq (4).

    Raises:
        KeyError: If *schedule_fn* is not a registered schedule.
    """
    if schedule_fn not in _DERIV_BY_FN:
        raise KeyError(
            "No analytic derivative registered for "
            f"{getattr(schedule_fn, '__name__', schedule_fn)!r}"
        )
    return _DERIV_BY_FN[schedule_fn]


def get_schedule(name: str) -> Callable[[Tensor], Tensor]:
    """Look up a noise schedule by name.

    Args:
        name: One of ``"linear"``, ``"cosine"``, ``"cosine_sq"``.

    Returns:
        The schedule function ``alpha(t)``.

    Raises:
        KeyError: If *name* is not registered.
    """
    if name not in _SCHEDULE_MAP:
        raise KeyError(
            f"Unknown schedule '{name}'. Available: {list(_SCHEDULE_MAP.keys())}"
        )
    return _SCHEDULE_MAP[name]