"""Minimal direct prompt-only H3 video sampler for visual smoke previews."""

import time

import torch

from .packing import H3PromptPacker, unpatchify_video


def shifted_sigma(base: torch.Tensor, shift: float) -> torch.Tensor:
    """H3 flow-SNR shift used by Comfy's ModelSamplingDiscreteFlow."""
    return shift * base / (1 + (shift - 1) * base)


def beta_sigmas(steps: int, *, device: torch.device | str, alpha: float = 0.6, beta: float = 0.6) -> torch.Tensor:
    """Comfy's discrete beta scheduler over H3's 1,000-entry shift-12 table."""
    from scipy.stats import beta as beta_distribution

    table = shifted_sigma(torch.arange(1, 1001, device=device, dtype=torch.float32) / 1000, 12.0)
    fractions = 1.0 - torch.arange(steps, device=device, dtype=torch.float64).cpu().numpy() / steps
    indices = torch.from_numpy((999 * beta_distribution.ppf(fractions, alpha, beta)).round().astype("int64")).to(device)
    indices = torch.unique_consecutive(indices)
    return torch.cat((table[indices], table.new_zeros(1)))


def res_multistep_update(x: torch.Tensor, denoised: torch.Tensor, sigma: torch.Tensor, sigma_down: torch.Tensor, old_denoised: torch.Tensor | None, old_sigma_down: torch.Tensor | None, previous_sigma: torch.Tensor | None) -> torch.Tensor:
    """Deterministic H3/Comfy RES multistep update."""
    if sigma_down == 0 or old_denoised is None:
        return x + (x - denoised) / sigma * (sigma_down - sigma)
    t, t_old, t_next, t_prev = sigma.log().neg(), old_sigma_down.log().neg(), sigma_down.log().neg(), previous_sigma.log().neg()
    h = t_next - t
    c2 = (t_prev - t_old) / h
    phi1 = torch.expm1(-h) / -h
    phi2 = (phi1 - 1.0) / -h
    b1 = torch.nan_to_num(phi1 - phi2 / c2, nan=0.0)
    b2 = torch.nan_to_num(phi2 / c2, nan=0.0)
    return torch.exp(-h) * x + h * (b1 * denoised + b2 * old_denoised)


def _unpack_audio(rows: torch.Tensor) -> torch.Tensor:
    steps = rows.shape[0] // 2
    return rows.reshape(2, steps, 32).permute(2, 0, 1).unsqueeze(0)


def _audio_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
    base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
    return 3.0 * base / (1.0 + (3.0 - 1.0) * base)


def _model_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
    """Comfy BaseModel's flow timestep round-trip seen by the H3 diffusion model."""
    return (video_sigma * 1000.0).float() / 1000.0


def _decode_audio_latent(audio_carried: torch.Tensor, *, shift_video: float = 12.0, shift_audio: float = 3.0) -> torch.Tensor:
    """Convert Comfy's AV sampler-carried audio state back to native audio VAE latents."""
    return audio_carried * (shift_audio / shift_video)


@torch.inference_mode()
def sample_video_res_multistep(
    model,
    packer: H3PromptPacker,
    text: torch.Tensor,
    video: torch.Tensor,
    audio: torch.Tensor,
    *,
    steps: int = 12,
    model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
    return_audio: bool = False,
) -> torch.Tensor:
    """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
    sigmas = beta_sigmas(steps, device=video.device)
    audio_carried = audio
    video_history = audio_history = None
    video_history_sigma = audio_history_sigma = None
    total_steps = len(sigmas) - 1
    started = time.perf_counter()
    for index, sigma in enumerate(sigmas[:-1], start=1):
        step_started = time.perf_counter()
        previous_index = index - 1
        sigma_down = sigmas[index]
        sigma_audio = _audio_sigma(sigma)
        carry = sigma_audio / sigma
        native_audio = audio_carried.to(torch.bfloat16) * carry
        step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
        hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps)
        raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
        raw_video = raw_video.to(torch.bfloat16).float()
        raw_audio = raw_audio.to(torch.bfloat16)
        velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
        velocity_audio = (
            (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
            + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
        ).float()
        video_denoised = video - sigma * velocity_video
        audio_denoised = audio_carried - sigma * velocity_audio
        previous_sigma = sigmas[previous_index - 1] if previous_index else None
        video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma)
        audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma)
        video_history, audio_history = video_denoised, audio_denoised
        video_history_sigma = audio_history_sigma = sigma_down
        elapsed = time.perf_counter() - started
        eta = elapsed / index * (total_steps - index)
        print(
            f"sampling step {index}/{total_steps}: "
            f"{time.perf_counter() - step_started:.1f}s, elapsed {elapsed:.1f}s, eta {eta:.1f}s",
            flush=True,
        )
    return (video, _decode_audio_latent(audio_carried)) if return_audio else video


@torch.inference_mode()
def sample_video_euler(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 2) -> torch.Tensor:
    """Use Euler updates to obtain a visual-only H3 preview, not parity sampling."""
    sigmas = beta_sigmas(steps, device=video.device)
    for index in range(steps):
        sigma = sigmas[index]
        hidden, timesteps, segments, positions, video_segment, audio_segment = packer(text, video, audio, _model_sigma(sigma))
        velocity, _ = model(hidden, timesteps, positions, segments, video_segment, audio_segment)
        velocity = unpatchify_video(velocity, video.shape[2], video.shape[-2], video.shape[-1])
        video.add_(velocity.to(video.dtype), alpha=float(sigmas[index + 1] - sigma))
    return video