"""Pull cross-attention scores and word confidence on the dosing audio.

Goal: see whether the AED's cross-attention distribution and per-word
confidence flag the hallucinated "5 milligramm 101" cluster that we
already identified via zero-duration word offsets.

Two new signals beyond what probe_canary_ts.py captured:
  1. word_confidence — per-word probability proxy. Hallucinations *should*
     have lower confidence than real words, though not guaranteed because
     AEDs hallucinate with high confidence too.
  2. xatt_scores — cross-attention from decoder tokens to encoder frames.
     For each output token we summarize where it "looked" in the audio:
       - peak_frame (argmax) and peak_time_s
       - entropy of the distribution (flat = no clear acoustic anchor)
     Hallucinated tokens should show high entropy (no clear acoustic peak).

Output:
  /sweep/probe_attn_conf.json — text + timestamp + word_confidence + per-token
    xatt summary (argmax frame, entropy)
"""
import json
import math
import subprocess
import tempfile

import torch
from omegaconf import OmegaConf

from nemo.collections.asr.models import EncDecMultiTaskModel
from nemo.collections.asr.models.aed_multitask_models import (
    MultiTaskTranscriptionConfig,
    parse_multitask_prompt,
)


WAV_SRC = "/sweep/case_c414cf52_15-01-39.m4a"
OUT_PATH = "/sweep/probe_attn_conf_kwargs_returnhyp.json"


def transcode(src: str) -> str:
    dst = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    dst.close()
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", src,
                    "-ac", "1", "-ar", "16000", "-f", "wav", dst.name],
                   check=True)
    return dst.name


def to_jsonable(o):
    if isinstance(o, torch.Tensor):
        return o.detach().cpu().tolist()
    if hasattr(o, "tolist"):
        return o.tolist()
    if hasattr(o, "item"):
        return o.item()
    return str(o)


def entropy_of(probs: torch.Tensor) -> float:
    p = probs.clamp_min(1e-12)
    return float(-(p * p.log()).sum().item())


def summarize_xatt(xatt_scores, frame_stride_s: float):
    """xatt_scores: List[Tensor] over decoder layers.

    Multi-head attention output shape is typically [heads, tokens, frames]
    or [tokens, frames] depending on the decoder. We mean-pool away any
    leading dimensions until we get a 2D [tokens, frames] matrix.
    """
    if not xatt_scores:
        return None
    layers = [x.float() for x in xatt_scores]
    print(f"  xatt layer shapes: {[tuple(l.shape) for l in layers]}", flush=True)
    avg = torch.stack(layers, dim=0).mean(dim=0)
    # Collapse all leading dims (batch, heads, ...) until 2D
    while avg.dim() > 2:
        avg = avg.mean(dim=0)
    n_tokens, n_frames = avg.shape
    summaries = []
    for t in range(n_tokens):
        row = avg[t]
        norm = row / row.sum().clamp_min(1e-12)
        peak = int(norm.argmax().item())
        summaries.append({
            "token_idx": t,
            "peak_frame": peak,
            "peak_time_s": round(peak * frame_stride_s, 3),
            "peak_prob": round(float(norm[peak].item()), 4),
            "entropy": round(entropy_of(norm), 4),
            "max_log_entropy": round(math.log(n_frames), 4),
        })
    return {"n_tokens": n_tokens, "n_frames": n_frames,
            "frame_stride_s": frame_stride_s,
            "n_layers": len(layers), "per_token": summaries}


def main():
    print("loading Canary 1B v2", flush=True)
    model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b-v2")
    model = model.to("cuda").to(torch.bfloat16)
    model.eval()

    # Patch the live decoding config:
    # - greedy strategy (confidence is only filled in greedy path, not beam)
    # - return_xattn_scores=True
    # - confidence_cfg with all preserve flags
    # - enable_chunking=False is set via override_config below (NeMo gates xatt
    #   on chunking; chunking also blocks word-level confidence aggregation).
    cfg = OmegaConf.to_container(model.cfg.decoding, resolve=True)
    cfg["strategy"] = "greedy"
    cfg["return_xattn_scores"] = True
    cfg.setdefault("confidence_cfg", {})
    cfg["confidence_cfg"]["preserve_word_confidence"] = True
    cfg["confidence_cfg"]["preserve_token_confidence"] = True
    cfg["confidence_cfg"]["preserve_frame_confidence"] = True
    cfg["confidence_cfg"]["aggregation"] = "min"  # most signal-preserving
    cfg.setdefault("greedy", {})
    cfg["greedy"]["preserve_token_confidence"] = True
    cfg["greedy"]["preserve_alignments"] = True
    print("applying decoding override:", json.dumps(cfg, indent=2), flush=True)
    model.change_decoding_strategy(OmegaConf.create(cfg))

    # Encoder frame stride: feature_stride * subsampling
    feature_stride = float(model.cfg.preprocessor["window_stride"])
    subsampling = 8  # FastConformer default; verified empirically against
                    # 80ms-per-frame in probe_tdt.md
    frame_stride_s = feature_stride * subsampling
    print(f"frame stride: {frame_stride_s} s", flush=True)

    wav = transcode(WAV_SRC)
    # Clean single-knob ablation: same override_config for both runs,
    # only `enable_chunking` differs. All other fields (batch_size, prompt,
    # timestamps, return_hypotheses) held constant.
    print(f"running transcribe(KWARGS path, return_hypotheses=True, timestamps=True) on {WAV_SRC}",
          flush=True)
    # Final clean test: kwargs path (like the original Mode A run that
    # showed None confidence), but with return_hypotheses=True forced.
    # If confidence appears, then `return_hypotheses` was the real
    # confounder all along — not timestamps, not the call form.
    with torch.amp.autocast("cuda", enabled=True, dtype=torch.bfloat16):
        with torch.no_grad():
            hyps = model.transcribe([wav], source_lang="de", target_lang="de",
                                     pnc="yes", timestamps=True, batch_size=1,
                                     return_hypotheses=True)
    h = hyps[0]

    out = {
        "text": getattr(h, "text", None),
        "score": getattr(h, "score", None),
        "timestamp": getattr(h, "timestamp", None),
        "tokens": getattr(h, "tokens", None),
        "word_confidence": getattr(h, "word_confidence", None),
        "token_confidence": getattr(h, "token_confidence", None),
        "xatt_summary": summarize_xatt(getattr(h, "xatt_scores", None),
                                        frame_stride_s),
    }

    with open(OUT_PATH, "w") as f:
        json.dump(out, f, default=to_jsonable, ensure_ascii=False, indent=2)

    # Quick sanity print
    print(f"\ntext: {h.text}", flush=True)
    print(f"score: {h.score}", flush=True)
    wc = getattr(h, "word_confidence", None)
    print(f"word_confidence: {len(wc) if wc is not None else 'None'} entries", flush=True)
    tc = getattr(h, "token_confidence", None)
    print(f"token_confidence: {len(tc) if tc is not None else 'None'} entries", flush=True)
    xs = getattr(h, "xatt_scores", None)
    if xs is not None:
        print(f"xatt_scores: {len(xs)} layer(s); shapes={[tuple(x.shape) for x in xs]}", flush=True)
    else:
        print("xatt_scores: None", flush=True)
    print(f"\nfull dump -> {OUT_PATH}", flush=True)


if __name__ == "__main__":
    main()