"""Trace when an H3 audio-boundary artifact emerges during base sampling."""

import argparse
import json
import math
from pathlib import Path

import torch

from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
from h3_blackwell_runtime.sampler import _decode_audio_latent, sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents


def dbfs(value: float) -> float:
    return 20.0 * math.log10(max(value, 1e-20))


def latent_metrics(latent: torch.Tensor) -> dict:
    frames = latent.float().movedim(-1, 0).flatten(1)
    return {
        "first_4_rms": float(frames[:4].square().mean().sqrt()),
        "frames_4_20_rms": float(frames[4:20].square().mean().sqrt()),
        "first_frame_rms": float(frames[0].square().mean().sqrt()),
        "frame_0_to_1_delta_rms": float((frames[1] - frames[0]).square().mean().sqrt()),
    }


def waveform_metrics(waveform: torch.Tensor) -> dict:
    waveform = waveform.float()
    first_100ms = waveform[..., :3200]
    first_500ms = waveform[..., :16000]
    derivative = (first_500ms[..., 1:] - first_500ms[..., :-1]).abs()
    return {
        "first_sample": waveform[..., 0].flatten().tolist(),
        "first_100ms_peak_dbfs": dbfs(float(first_100ms.abs().max())),
        "first_100ms_rms_dbfs": dbfs(float(first_100ms.square().mean().sqrt())),
        "first_500ms_peak_dbfs": dbfs(float(first_500ms.abs().max())),
        "first_500ms_rms_dbfs": dbfs(float(first_500ms.square().mean().sqrt())),
        "largest_derivative": float(derivative.max()),
    }


parser = argparse.ArgumentParser()
parser.add_argument("--benchmark", type=Path, required=True)
parser.add_argument("--trace", type=Path, required=True)
parser.add_argument("--report", type=Path, required=True)
parser.add_argument("--attention", default="sage2")
args = parser.parse_args()

benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention))
video, initial_audio, aligned_frames = random_av_latents(
    benchmark["resolution"][0],
    benchmark["resolution"][1],
    benchmark["frames"],
    benchmark["seed"],
    device=runtime.config.device,
)
text = runtime.refiner(runtime.conditioner(benchmark["prompt"]))
trace = []
video, final_audio = sample_video_res_multistep(
    runtime.model,
    runtime.packer,
    text,
    video,
    initial_audio,
    steps=benchmark["steps"],
    seed=benchmark["seed"],
    return_audio=True,
    audio_step_trace=trace,
)

report = {
    "benchmark": str(args.benchmark),
    "attention": args.attention,
    "seed": benchmark["seed"],
    "frames": aligned_frames,
    "steps": [],
}
with torch.inference_mode():
    for entry in trace:
        denoised = _decode_audio_latent(entry["audio_denoised"]).to(
            "cuda", dtype=next(runtime.audio_vae.parameters()).dtype,
        )
        waveform = runtime.audio_vae.decode(denoised).cpu()[0]
        report["steps"].append({
            "step": entry["step"],
            "video_sigma": entry["video_sigma"],
            "audio_sigma": entry["audio_sigma"],
            "latent": latent_metrics(denoised.cpu()),
            "denoised_waveform": waveform_metrics(waveform),
        })

args.trace.parent.mkdir(parents=True, exist_ok=True)
torch.save({
    "initial_audio": initial_audio.detach().cpu(),
    "final_audio": final_audio.detach().cpu(),
    "steps": trace,
    "prompt": benchmark["prompt"],
    "seed": benchmark["seed"],
}, args.trace)
args.report.parent.mkdir(parents=True, exist_ok=True)
serialized = json.dumps(report, indent=2)
args.report.write_text(serialized + "\n", encoding="utf-8")
print(serialized)