import torch

# ---------- diffusion schedule (pre-computed once) ----------
_NUM_STEPS = 1000
# βₜ: from 1e-5 → 5e-3 on a sigmoid scale
_betas = torch.sigmoid(torch.linspace(-6, 6, _NUM_STEPS)) * (5e-3 - 1e-5) + 1e-5
_alphas = 1.0 - _betas
_alphas_cumprod = torch.cumprod(_alphas, dim=0)           # ᾱ_t
_sqrt_alphas_cumprod = _alphas_cumprod.sqrt()             # √ᾱ_t
_sqrt_one_minus_alphas_cumprod = (1.0 - _alphas_cumprod).sqrt()  # √(1-ᾱ_t)

# ---------- noise helper ----------
def add_diffusion_noise(x: torch.Tensor, step: int) -> torch.Tensor:
    """
    Forward-process diffusion noise at a given step.

    Args:
        x   : Tensor (C,H,W) or (B,C,H,W), range [-1,1] 或 [0,1]
        step: int in [0, 999]

    Returns:
        Noisy tensor with same shape / dtype / device
    """
    if not (0 <= step < _NUM_STEPS):
        raise ValueError(f"`step` must be in [0, {_NUM_STEPS-1}], got {step}")

    # move schedule scalars to correct device / dtype
    device, dtype = x.device, x.dtype
    c1 = _sqrt_alphas_cumprod[step].to(device=device, dtype=dtype)
    c2 = _sqrt_one_minus_alphas_cumprod[step].to(device=device, dtype=dtype)

    noise = torch.randn_like(x)
    return c1 * x + c2 * noise
