"""Validate a trained draft checkpoint: does it produce coherent output?

Runs the draft standalone (greedy generation) on a single audio clip and compares
to the target's output. Also runs one prefill of both on the same audio to
measure top-1 argmax agreement (a proxy for acceptance rate).

Everything runs on CPU (audio encoder can't run on trn2, and for validation
CPU is fine).
"""

from __future__ import annotations

import argparse
import time
import torch
from pathlib import Path
from transformers import VoxtralForConditionalGeneration, AutoProcessor


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--target", default="mistralai/Voxtral-Mini-3B-2507")
    p.add_argument("--draft", required=True, help="Path to trained draft checkpoint dir")
    p.add_argument("--audio", default="/mnt/data/LibriSpeech/dev-clean/2035/147960/2035-147960-0015.flac")
    p.add_argument("--max-new-tokens", type=int, default=64)
    p.add_argument("--language", default="en")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    torch.set_grad_enabled(False)
    dtype = torch.bfloat16

    print(f"[validate] Loading target {args.target}")
    target = VoxtralForConditionalGeneration.from_pretrained(
        args.target, torch_dtype=dtype, low_cpu_mem_usage=True,
    ).eval()
    proc = AutoProcessor.from_pretrained(args.target)

    print(f"[validate] Loading draft {args.draft}")
    draft = VoxtralForConditionalGeneration.from_pretrained(
        args.draft, torch_dtype=dtype, low_cpu_mem_usage=True,
    ).eval()
    print(f"[validate] Target layers: {target.config.text_config.num_hidden_layers}")
    print(f"[validate] Draft layers:  {draft.config.text_config.num_hidden_layers}")

    inputs = proc.apply_transcription_request(
        language=args.language, audio=args.audio, model_id=args.target,
    )

    # -- Target greedy
    print("\n[validate] Target greedy generation")
    t0 = time.perf_counter()
    out_t = target.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False,
                            temperature=None, top_p=None)
    gen_t = out_t[0, inputs["input_ids"].shape[1]:]
    target_text = proc.tokenizer.decode(gen_t, skip_special_tokens=True).strip()
    print(f"[validate] Target ({time.perf_counter()-t0:.1f}s, {gen_t.shape[0]} tok):")
    print(f"           {target_text!r}")

    # -- Draft greedy (standalone)
    print("\n[validate] Draft greedy generation (standalone)")
    t0 = time.perf_counter()
    out_d = draft.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False,
                           temperature=None, top_p=None)
    gen_d = out_d[0, inputs["input_ids"].shape[1]:]
    draft_text = proc.tokenizer.decode(gen_d, skip_special_tokens=True).strip()
    print(f"[validate] Draft ({time.perf_counter()-t0:.1f}s, {gen_d.shape[0]} tok):")
    print(f"           {draft_text!r}")
    print(f"[validate] First 20 raw tokens: {gen_d[:20].tolist()}")

    # -- Top-1 agreement: how often does draft's argmax match target's argmax
    # given TEACHER-FORCED context = target's greedy sequence?
    print("\n[validate] Top-1 agreement (teacher-forced)")
    # Build full sequence: prefix + target's generated tokens
    full_ids = out_t.clone()  # [1, prefix_len + n_gen_target]
    prefix_len = inputs["input_ids"].shape[1]
    # Forward through target (get logits at every position) and same for draft
    with torch.no_grad():
        target_out = target(input_ids=full_ids, input_features=inputs["input_features"])
        draft_out  = draft(input_ids=full_ids,  input_features=inputs["input_features"])
    target_argmax = target_out.logits.argmax(-1)  # [1, seq_len]
    draft_argmax  = draft_out.logits.argmax(-1)
    # Agreement over positions [prefix_len - 1 ... end - 1] (positions that predict target tokens)
    agree_positions = target_argmax[0, prefix_len-1:-1] == draft_argmax[0, prefix_len-1:-1]
    total_positions = agree_positions.numel()
    n_agree = agree_positions.sum().item()
    print(f"[validate] Top-1 agreement: {n_agree}/{total_positions} = {n_agree/total_positions*100:.1f}%")

    # -- Also print divergences
    if total_positions > 0:
        print(f"[validate] First 20 positions (target vs draft argmax):")
        for i in range(min(20, total_positions)):
            tgt_id = target_argmax[0, prefix_len-1+i].item()
            dft_id = draft_argmax[0, prefix_len-1+i].item()
            match = "✓" if tgt_id == dft_id else "✗"
            tgt_tok = proc.tokenizer.decode([tgt_id])
            dft_tok = proc.tokenizer.decode([dft_id])
            print(f"           pos {i:3d}  {match}  tgt={tgt_id} ({tgt_tok!r})  dft={dft_id} ({dft_tok!r})")

    print(f"\n[validate] Summary:")
    print(f"           draft coherent: {'yes' if not draft_text.startswith(('\",\"', '.,.', ',,,')) else 'no (gibberish)'}")
    print(f"           byte-identical to target: {draft_text == target_text}")
    print(f"           top-1 agreement:          {n_agree/total_positions*100:.1f}%")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())