"""The paper's DDPM sampler (model.sample_ddpm) with two additions for the demo.

- Guidance, as the old listener had it: a second pass with the encoded grid zeroed (its
  ``force_uncond``), eps = eps_uncond + w * (eps_cond - eps_uncond). Following Kynkäänniemi et al.
  (NeurIPS 2024), it can be limited to an interval of noise levels; w = 1 is the paper's sampler.
- Frozen frames: latent frames given in ``known`` are replaced at every step by the known latent
  noised to that step's level (RePaint-style; "frozen overlap" continuation), so a chunk can be
  sampled as the continuation of the previous one.

With w = 1 and nothing known this is exactly model.sample_ddpm.
"""
from __future__ import annotations

from typing import Any, Mapping

import torch

from model import _prepare_batch_tensors, apply_seq_mask, clip_x0, resolve_inference_geometry


@torch.no_grad()
def sample(
    *,
    model: Any,
    diffusion: Any,
    batch: Mapping[str, Any],
    device: torch.device,
    sample_seed: int,
    x0_clip_bound: float,
    target_token_rate_hz: float,
    inference_num_beats: int = 4,
    inference_geometry: Mapping[str, Any] | None = None,
    guidance: float = 1.0,
    guidance_steps: tuple[int, int] | None = None,
    known: torch.Tensor | None = None,
    known_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Returns the normalized latent [B,T,D]. ``guidance_steps`` = (lo, hi) applies guidance only for
    lo <= step <= hi (steps count down from num_steps - 1). ``known``/``known_mask``: [B,T,D] / [B,T]."""
    prepared = _prepare_batch_tensors(batch, device, require_target=False, require_timing=False)
    geometry = inference_geometry or resolve_inference_geometry(
        prepared,
        use_bpm_inference_geometry=True,
        inference_num_beats=int(inference_num_beats),
        target_token_rate_hz=float(target_token_rate_hz),
    )
    target_mask = geometry["target_valid_mask_bt"]
    cond, cond_mask = model.encode_conditioning(
        grid=prepared["grid"],
        grid_ids=prepared["grid_ids"],
        grid_times_sec=geometry["grid_times_sec"],
        token_times_sec=geometry["token_times_sec"],
        target_valid_mask_bt=target_mask,
        grid_valid_mask_bt=prepared["grid_valid_mask"],
    )
    uncond = torch.zeros_like(cond)
    gen = torch.Generator(device=device).manual_seed(int(sample_seed))
    x = apply_seq_mask(
        torch.randn(target_mask.shape[0], target_mask.shape[1], int(model.cfg.x_dim), device=device, generator=gen),
        target_mask,
    )
    keep = known_mask[..., None].to(device) if known is not None else None

    def eps_at(x_t: torch.Tensor, t: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        return model(
            x_t=x_t,
            t=t,
            target_valid_mask_bt=target_mask,
            token_times_sec=geometry["token_times_sec"],
            cond_btd=c,
            cond_valid_mask_bt=cond_mask,
        )

    for step in reversed(range(diffusion.num_steps)):
        t = torch.full((x.shape[0],), step, device=device, dtype=torch.long)
        if keep is not None:  # the known frames at this step's noise level
            x = torch.where(keep, diffusion.q_sample(known.to(device), t, torch.randn(x.shape, device=device, generator=gen)), x)
        eps = eps_at(x, t, cond)
        guided = guidance != 1.0 and (guidance_steps is None or guidance_steps[0] <= step <= guidance_steps[1])
        if guided:
            eps_u = eps_at(x, t, uncond)
            eps = eps_u + float(guidance) * (eps - eps_u)
        x0_hat = apply_seq_mask(clip_x0(diffusion.predict_x0_from_eps(x, t, eps), x0_clip_bound), target_mask)
        mean = diffusion.posterior_mean_from_x0(x, t, x0_hat)
        if step > 0:
            z = torch.randn(x.shape, dtype=x.dtype, device=device, generator=gen)
            x = mean + torch.sqrt(diffusion.posterior_variance[t].view(-1, 1, 1)) * z
        else:
            x = mean
        x = apply_seq_mask(x, target_mask)
    if keep is not None:
        x = torch.where(keep, known.to(device), x)
    return x