#!/usr/bin/env python3
"""Verify the exported ONNX graphs against the original torch model.

  1. Encoder parity : ONNX encoder vs torch EncoderForExport (max abs diff).
  2. Decoder parity : ONNX decoder step-0 logits vs torch DecoderMergedForExport.
  3. End-to-end     : greedy decode via ONNX (Rust prompt) -> text; compare to model.generate().

Usage: python verify.py [sample.wav] [encoder_suffix] [decoder_suffix]
       python verify.py sample_en.wav ""      ""        # fp32
       python verify.py sample_en.wav "_q4"   "_q4"     # q4
"""
import sys
import numpy as np
import onnxruntime as ort
import torch

from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
from export_merged import EncoderForExport, DecoderMergedForExport, NUM_DECODER_LAYERS

MODEL_DIR = "model_ar"
WAV = sys.argv[1] if len(sys.argv) > 1 else "sample_en.wav"
ENC_SFX = sys.argv[2] if len(sys.argv) > 2 else ""
DEC_SFX = sys.argv[3] if len(sys.argv) > 3 else ""

# Rust build_prompt(language="ar", punctuation=True) -> token ids
PROMPT_TOKENS = ["▁", "<|startofcontext|>", "<|startoftranscript|>", "<|emo:undefined|>",
                 "<|ar|>", "<|ar|>", "<|pnc|>", "<|noitn|>", "<|notimestamp|>", "<|nodiarize|>"]
EOS_ID = 3
H, D = 8, 128


def sess(path):
    return ort.InferenceSession(path, providers=["CPUExecutionProvider"])


def main():
    import soundfile as sf
    print(f"Loading torch model + processor from {MODEL_DIR}")
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        MODEL_DIR, trust_remote_code=True, dtype=torch.float32).eval()
    proc = AutoProcessor.from_pretrained(MODEL_DIR, trust_remote_code=True)

    vocab = {t: i for t, i in __import__("json").load(
        open(f"{MODEL_DIR}/tokenizer.json", encoding="utf-8"))["model"]["vocab"].items()}
    id2tok = {i: t for t, i in vocab.items()}
    prompt_ids = [vocab[t] for t in PROMPT_TOKENS]
    print("prompt ids:", prompt_ids)

    # Build the featurizer from config.preprocessor (preprocessor_config.json lacks params, so the
    # AutoProcessor falls back to nfilt=64 — wrong; the real params live in config.json).
    FE = type(proc.feature_extractor)  # CohereAsrFeatureExtractor (already trust_remote_code-loaded)
    pp = model.config.preprocessor
    fe = FE(
        feature_size=pp["features"], sampling_rate=pp["sample_rate"],
        n_window_size=int(pp["window_size"] * pp["sample_rate"]),
        n_window_stride=int(pp["window_stride"] * pp["sample_rate"]),
        window=pp["window"], normalize=pp["normalize"], n_fft=pp["n_fft"],
        log=pp["log"], dither=pp["dither"], pad_to=pp["pad_to"], mel_norm="slaney",
    )

    wav, sr = sf.read(WAV)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    wav = wav.astype(np.float32)
    feats = fe(wav, sampling_rate=16000, return_tensors="pt")
    input_features = feats["input_features"]  # (B, 128, T) -> (B, T, 128)
    if input_features.shape[-1] != 128:
        input_features = input_features.transpose(1, 2)
    print("input_features:", tuple(input_features.shape))

    # ---- 1. encoder parity ----
    with torch.no_grad():
        torch_enc = EncoderForExport(model).eval()(input_features)
    enc_sess = sess(f"onnx/encoder_model{ENC_SFX}.onnx")
    onnx_enc = enc_sess.run(["last_hidden_state"],
                            {"input_features": input_features.numpy()})[0]
    diff = np.abs(torch_enc.numpy() - onnx_enc).max()
    print(f"[encoder] shape {onnx_enc.shape}  max|torch-onnx| = {diff:.3e}")

    # ---- 2. decoder step-0 parity ----
    dec_sess = sess(f"onnx/decoder_model_merged{DEC_SFX}.onnx")
    P = len(prompt_ids)
    empties = {}
    feeds = {
        "input_ids": np.array([prompt_ids], np.int64),
        "attention_mask": np.ones((1, P), np.int64),
        "position_ids": np.arange(P, dtype=np.int64)[None],
        "num_logits_to_keep": np.array(1, np.int64),
        "encoder_hidden_states": onnx_enc,
    }
    for i in range(NUM_DECODER_LAYERS):
        for kind in ("decoder", "encoder"):
            for kv in ("key", "value"):
                feeds[f"past_key_values.{i}.{kind}.{kv}"] = np.zeros((1, H, 0, D), np.float32)
    onnx_logits = dec_sess.run(["logits"], feeds)[0]
    with torch.no_grad():
        args = (torch.tensor([prompt_ids]), torch.ones(1, P, dtype=torch.long),
                torch.arange(P)[None], torch.tensor(1),
                torch.tensor(onnx_enc))
        past = []
        for _ in range(NUM_DECODER_LAYERS):
            past += [torch.zeros(1, H, 0, D), torch.zeros(1, H, 0, D),
                     torch.zeros(1, H, 0, D), torch.zeros(1, H, 0, D)]
        torch_logits = DecoderMergedForExport(model).eval()(*args, *past)[0]
    dl = np.abs(torch_logits.numpy() - onnx_logits).max()
    print(f"[decoder] logits {onnx_logits.shape}  max|torch-onnx| = {dl:.3e}")

    # ---- 3. end-to-end greedy decode via ONNX ----
    def decode_text(ids):
        out, buf = "", bytearray()
        for tid in ids:
            tok = id2tok.get(tid, "")
            if tok.startswith("<0x") and tok.endswith(">"):
                buf.append(int(tok[3:-1], 16)); continue
            if buf:
                out += buf.decode("utf-8", "replace"); buf = bytearray()
            if tok.startswith("<|") or tok in ("<s>", "</s>"):
                continue
            out += tok.replace("▁", " ")
        if buf:
            out += buf.decode("utf-8", "replace")
        return out.strip()

    past = {f"past_key_values.{i}.{k}.{kv}": np.zeros((1, H, 0, D), np.float32)
            for i in range(NUM_DECODER_LAYERS) for k in ("decoder", "encoder") for kv in ("key", "value")}
    cur = prompt_ids[:]
    gen = []
    for step in range(448):
        L = len(cur)
        attn = len(prompt_ids) + step
        pos = (np.arange(L) if step == 0 else np.array([attn - 1])).astype(np.int64)[None]
        feeds = {"input_ids": np.array([cur], np.int64),
                 "attention_mask": np.ones((1, attn if step else L), np.int64),
                 "position_ids": pos, "num_logits_to_keep": np.array(1, np.int64),
                 "encoder_hidden_states": onnx_enc, **past}
        outs = dec_sess.run(None, feeds)
        names = [o.name for o in dec_sess.get_outputs()]
        res = dict(zip(names, outs))
        nxt = int(res["logits"][0, -1].argmax())
        if nxt == EOS_ID:
            break
        gen.append(nxt)
        for i in range(NUM_DECODER_LAYERS):
            for k in ("decoder", "encoder"):
                for kv in ("key", "value"):
                    v = res[f"present.{i}.{k}.{kv}"]
                    if v.shape[2] != 0:
                        past[f"past_key_values.{i}.{k}.{kv}"] = v
        cur = [nxt]
    print(f"\n[ONNX greedy] ({len(gen)} tok): {decode_text(gen)!r}")

    # ---- reference: model.generate ----
    try:
        with torch.no_grad():
            g = model.generate(input_features=feats["input_features"], max_new_tokens=200)
        txt = proc.batch_decode(g, skip_special_tokens=True)[0] if hasattr(proc, "batch_decode") \
            else proc.tokenizer.decode(g[0], skip_special_tokens=True)
        print("[generate]  ", txt.strip()[:200])
    except Exception as e:
        print("[generate] skipped:", type(e).__name__, str(e)[:160])


if __name__ == "__main__":
    main()