#!/usr/bin/env python3
"""Cossim debug : C++ omnivoice-tts vs Python OmniVoice on voice cloning.

Inputs (relative to CWD = tests/) :
    ../examples/prompt.txt       target text fed to both pipelines
    ../examples/freeman.txt      transcript of the cloning reference
    ../examples/freeman.wav      cloning reference audio, any rate, any layout

Both sides run with seed=42, F32 weights, language=English, no pre or post
process. The reference audio is resampled to 24 kHz mono inside both
pipelines.

Dumps land in cpp/ (C++) and python/ (Python) and are compared pair by
pair. All paths are relative.
"""

import argparse
import os
import struct
import subprocess
import sys

# Strict F32 matmul on both sides. NVIDIA_TF32_OVERRIDE=0 forces full FP32
# mantissa in cuBLAS for both PyTorch and the C++ child via inheritance.
# Must be set BEFORE torch imports so the cuBLAS handle reads it on init.
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
import torchaudio

# Belt and suspenders : disable PyTorch's own TF32 toggles too. Some code
# paths bypass NVIDIA_TF32_OVERRIDE through cudnn or torch internal flags.
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32       = False

from omnivoice import OmniVoice
from omnivoice.utils.common import fix_random_seed

BIN        = "../build/omnivoice-tts"
MODEL_LM_T = "../models/omnivoice-base-{q}.gguf"
MODEL_CDC_T = "../models/omnivoice-tokenizer-{q}.gguf"
CKPT       = "../checkpoints/OmniVoice"
DUMP_CPP   = "cpp"
DUMP_PT    = "python"

def ensure_dir(path):
    os.makedirs(path, exist_ok=True)

def save_dump(path, data):
    if isinstance(data, torch.Tensor):
        data = data.detach().to(torch.float32).cpu().numpy()
    data  = np.ascontiguousarray(data.astype(np.float32))
    shape = data.shape
    with open(path, "wb") as f:
        f.write(struct.pack("i", len(shape)))
        for s in shape:
            f.write(struct.pack("i", s))
        f.write(data.tobytes())

def load_dump(path):
    raw   = np.fromfile(path, dtype=np.uint8)
    ndim  = int(np.frombuffer(raw[0:4], dtype=np.int32)[0])
    shape = tuple(int(x) for x in np.frombuffer(raw[4:4 + 4 * ndim], dtype=np.int32))
    body  = np.frombuffer(raw[4 + 4 * ndim:], dtype=np.float32)
    return body.reshape(shape), shape

def cos(a, b):
    a = a.astype(np.float64).ravel()
    b = b.astype(np.float64).ravel()
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]
    # log_probs contains -inf at the audio_mask_id slot. Mask these symmetric
    # positions to 0 on both sides so they cancel out of the inner product.
    bad = ~(np.isfinite(a) & np.isfinite(b))
    if bad.any():
        a = np.where(bad, 0.0, a)
        b = np.where(bad, 0.0, b)
    d = float(np.linalg.norm(a) * np.linalg.norm(b))
    return float(np.dot(a, b) / d) if d > 1e-10 else 0.0

def stft_cos(a, b, win=2048, hop=512):
    # STFT magnitude cosine. Drops phase, so a constant time shift between
    # the two waveforms does not collapse the score. The plain cos() on
    # raw samples falls to ~0 the moment chunks land a few samples apart.
    a = a.astype(np.float64).ravel()
    b = b.astype(np.float64).ravel()
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]
    window = np.hanning(win)
    frames = (n - win) // hop + 1
    if frames <= 0:
        return 0.0
    sa = np.zeros((frames, win // 2 + 1))
    sb = np.zeros((frames, win // 2 + 1))
    for i in range(frames):
        s = i * hop
        sa[i] = np.abs(np.fft.rfft(a[s:s + win] * window))
        sb[i] = np.abs(np.fft.rfft(b[s:s + win] * window))
    return cos(sa.ravel(), sb.ravel())

def install_hooks(model, dump_dir):
    # Capture the raw input_ids row k=0 for cond and uncond at the first
    # forward pass. Lets a divergence in prompt construction (style, text,
    # ref_audio tokens) localize before the LM runs.
    seen_embed = {"done": False}
    orig_prepare = model._prepare_embed_inputs
    def hooked_prepare(input_ids, audio_mask):
        out = orig_prepare(input_ids, audio_mask)
        if not seen_embed["done"] and input_ids.dim() == 3 and input_ids.shape[0] >= 2:
            cond_ids   = input_ids[0, 0, :].detach().to(torch.float32).cpu().numpy()
            uncond_ids = input_ids[1, 0, :].detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, "prompt-cond-ids.bin"),   cond_ids)
            save_dump(os.path.join(dump_dir, "prompt-uncond-ids.bin"), uncond_ids)
            seen_embed["done"] = True
        return out
    model._prepare_embed_inputs = hooked_prepare

    seen = {"step0": False, "mg_tokens": False, "audio": False}
    orig_pred = model._predict_tokens_with_scoring
    def hooked_pred(c_logits, u_logits, gen_config):
        is_first = not seen["step0"]
        if is_first:
            c = c_logits.detach().to(torch.float32).cpu().numpy()
            u = u_logits.detach().to(torch.float32).cpu().numpy()
            if c.ndim == 4:
                c = c[0]
            if u.ndim == 4:
                u = u[0]
            save_dump(os.path.join(dump_dir, "lm-logits-step0-cond.bin"),   c)
            save_dump(os.path.join(dump_dir, "lm-logits-step0-uncond.bin"), u)
        pred_tokens, scores = orig_pred(c_logits, u_logits, gen_config)
        if is_first:
            # Reconstruct log_probs with the same op sequence as the Python
            # function so the dump is the actual log_probs the model used.
            with torch.no_grad():
                if gen_config.guidance_scale != 0:
                    cl = F.log_softmax(c_logits, dim=-1)
                    ul = F.log_softmax(u_logits, dim=-1)
                    lp = torch.log_softmax(cl + gen_config.guidance_scale * (cl - ul), dim=-1)
                else:
                    lp = F.log_softmax(c_logits, dim=-1)
                lp[..., model.config.audio_mask_id] = -float("inf")
            lp_arr = lp.detach().to(torch.float32).cpu().numpy()
            if lp_arr.ndim == 4:
                lp_arr = lp_arr[0]
            save_dump(os.path.join(dump_dir, "mg-log-probs-step0.bin"), lp_arr)

            pt_arr = pred_tokens.detach().to(torch.float32).cpu().numpy()
            if pt_arr.ndim == 3:
                pt_arr = pt_arr[0]
            save_dump(os.path.join(dump_dir, "mg-pred-tokens-step0.bin"), pt_arr)

            sc_arr = scores.detach().to(torch.float32).cpu().numpy()
            if sc_arr.ndim == 3:
                sc_arr = sc_arr[0]
            save_dump(os.path.join(dump_dir, "mg-scores-step0.bin"), sc_arr)
            seen["step0"] = True
        return pred_tokens, scores
    model._predict_tokens_with_scoring = hooked_pred

    orig_generate = model._generate_iterative
    def hooked_generate(task, gen_config):
        out = orig_generate(task, gen_config)
        if not seen["mg_tokens"]:
            save_dump(os.path.join(dump_dir, "mg-tokens.bin"), out[0])
            seen["mg_tokens"] = True
        return out
    model._generate_iterative = hooked_generate

    orig_decode = model.audio_tokenizer.decode
    def hooked_decode(*args, **kwargs):
        out = orig_decode(*args, **kwargs)
        if seen["audio"]:
            return out
        wav = getattr(out, "audio_values", out)
        if isinstance(wav, torch.Tensor):
            arr = wav.detach().to(torch.float32).cpu().numpy()
        else:
            arr = np.asarray(wav, dtype=np.float32)
        if arr.ndim == 3:
            arr = arr[0, 0]
        elif arr.ndim == 2:
            arr = arr[0]
        save_dump(os.path.join(dump_dir, "output-audio.bin"), arr)
        seen["audio"] = True
        return out
    model.audio_tokenizer.decode = hooked_decode

    # Capture the full reference audio code matrix [K, T_ref] before it goes
    # into the prompt, so a divergence in the codec encoder is visible at the
    # source rather than only via prompt-cond-ids row k=0.
    orig_encode = model.audio_tokenizer.encode
    seen_enc = {"done": False}
    def hooked_encode(*args, **kwargs):
        out = orig_encode(*args, **kwargs)
        if not seen_enc["done"]:
            codes = out.audio_codes if hasattr(out, "audio_codes") else out
            if isinstance(codes, torch.Tensor):
                arr = codes.detach().to(torch.float32).cpu().numpy()
                if arr.ndim == 3:
                    arr = arr[0]
                save_dump(os.path.join(dump_dir, "ref-audio-codes.bin"), arr)
                seen_enc["done"] = True
        return out
    model.audio_tokenizer.encode = hooked_encode

    # Capture the 16 kHz post resample audio and the HuBERT semantic features
    # at the boundary inside the audio tokenizer. Lets us bisect the codec
    # encoder : raw 24k -> resample 16k -> hubert -> sem_enc -> RVQ -> codes.
    orig_extract = model.audio_tokenizer._extract_semantic_features
    seen_sem = {"done": False}
    def hooked_extract(input_values):
        if not seen_sem["done"]:
            cfg = model.audio_tokenizer.config
            if cfg.sample_rate != cfg.semantic_sample_rate:
                v16 = torchaudio.functional.resample(input_values, cfg.sample_rate, cfg.semantic_sample_rate)
            else:
                v16 = input_values
            arr16 = v16[:, 0, :].squeeze(0).detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, "ref-audio-16k.bin"), arr16)
        out = orig_extract(input_values)
        if not seen_sem["done"]:
            feat = out.squeeze(0).detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, "ref-hubert-features.bin"), feat)
            seen_sem["done"] = True
        return out
    model.audio_tokenizer._extract_semantic_features = hooked_extract

    # Bisect HuBERT internals. The C++ counterpart marks the same 9 stages as
    # graph outputs in pipeline_codec_hubert_features_test and dumps them post
    # compute. Names match exactly so the Python compare loop can pair them.
    sm = model.audio_tokenizer.semantic_model
    seen_hub = {"feat": False, "proj_ln": False, "proj": False, "init": False,
                "l0":   False, "l5":      False, "l7":   False, "l9":   False, "l11":  False}

    def hub_hook_feat(module, inputs, output):
        # HF feature_extractor returns (B, C=512, T_feat). Squeeze gives
        # (C, T_feat) which matches the C++ ne=(T, C) tensor whose linear
        # buffer is laid out slow-first as (C, T). No transpose needed.
        if seen_hub["feat"]:
            return
        arr = output[0].detach().to(torch.float32).cpu().numpy()
        save_dump(os.path.join(dump_dir, "hubert-feat-extract.bin"), arr)
        seen_hub["feat"] = True
    sm.feature_extractor.register_forward_hook(hub_hook_feat)

    # feature_projection internal split : layer_norm first, then projection
    # Linear. The C++ side dumps the post LN tensor before the Linear, so we
    # mirror with a forward hook on the LN sub-module.
    def hub_hook_proj_ln(module, inputs, output):
        if seen_hub["proj_ln"]:
            return
        out_t = output[0] if isinstance(output, tuple) else output
        arr = out_t[0].detach().to(torch.float32).cpu().numpy()
        save_dump(os.path.join(dump_dir, "hubert-feat-proj-ln.bin"), arr)
        seen_hub["proj_ln"] = True
    sm.feature_projection.layer_norm.register_forward_hook(hub_hook_proj_ln)

    def hub_hook_proj(module, inputs, output):
        # feature_projection returns (B, T_feat, 768) already T-first.
        if seen_hub["proj"]:
            return
        out_t = output[0] if isinstance(output, tuple) else output
        arr = out_t[0].detach().to(torch.float32).cpu().numpy()
        save_dump(os.path.join(dump_dir, "hubert-feat-proj.bin"), arr)
        seen_hub["proj"] = True
    sm.feature_projection.register_forward_hook(hub_hook_proj)

    # encoder.layer_norm is the LN that follows pos_conv_embed add. Its output
    # equals the C++ enc_init output : x + pos_conv(x) -> LN.
    def hub_hook_init(module, inputs, output):
        if seen_hub["init"]:
            return
        arr = output[0].detach().to(torch.float32).cpu().numpy()
        save_dump(os.path.join(dump_dir, "hubert-enc-init.bin"), arr)
        seen_hub["init"] = True
    sm.encoder.layer_norm.register_forward_hook(hub_hook_init)

    # Mid stack taps : layer 0, 5, 7, 9, 11. Mirror the C++ states[1], [6],
    # [8], [10], [12] indexing (states[0] is enc-init dumped above). l7 and
    # l9 give the resolution to localize the explosion seen between l5 and l11.
    def make_layer_tap(key, fname):
        def hook(module, inputs, output):
            if seen_hub[key]:
                return
            h = output[0] if isinstance(output, tuple) else output
            arr = h[0].detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, fname), arr)
            seen_hub[key] = True
        return hook
    sm.encoder.layers[0].register_forward_hook(make_layer_tap("l0",  "hubert-l0.bin"))
    sm.encoder.layers[5].register_forward_hook(make_layer_tap("l5",  "hubert-l5.bin"))
    sm.encoder.layers[7].register_forward_hook(make_layer_tap("l7",  "hubert-l7.bin"))
    sm.encoder.layers[9].register_forward_hook(make_layer_tap("l9",  "hubert-l9.bin"))
    sm.encoder.layers[11].register_forward_hook(make_layer_tap("l11", "hubert-l11.bin"))

    # Bisect the LM forward by dumping per layer hidden states at step 0.
    # The C++ pipeline dumps cond and uncond at layers 0, 6, 13, 20 and the
    # final norm. Mirror that exactly so a per layer max_abs_diff comparison
    # localizes where the drift starts.
    seen_lm = {"done": False}
    def make_layer_hook(name):
        def hook(module, inputs, output):
            if seen_lm["done"]:
                return
            h = output[0] if isinstance(output, tuple) else output
            if not isinstance(h, torch.Tensor) or h.ndim < 3:
                return
            cond_arr = h[0].detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, f"lm-hidden-step0-cond-{name}.bin" if name else "lm-hidden-step0-cond.bin"), cond_arr)
            if h.shape[0] >= 2:
                unc_arr = h[1].detach().to(torch.float32).cpu().numpy()
                save_dump(os.path.join(dump_dir, f"lm-hidden-step0-uncond-{name}.bin" if name else "lm-hidden-step0-uncond.bin"), unc_arr)
        return hook
    lm_hook_handles = []
    for idx, name in [(0, "l0"), (1, "l1"), (2, "l2"), (3, "l3"), (4, "l4"), (5, "l5"), (6, "l6"),
                      (13, "l13"), (14, "l14"), (15, "l15"), (16, "l16"), (17, "l17"), (18, "l18"),
                      (19, "l19"), (20, "l20")]:
        if idx < len(model.llm.layers):
            lm_hook_handles.append(model.llm.layers[idx].register_forward_hook(make_layer_hook(name)))
    lm_hook_handles.append(model.llm.norm.register_forward_hook(make_layer_hook("")))
    # Trip seen_lm["done"] right after the first LM forward by hooking the
    # outermost forward via the existing _predict_tokens_with_scoring path
    # already gated on step0. We instead use an extra one-shot hook on the
    # last collected layer (norm) to flip the flag after its dump fires.
    def flip_done(module, inputs, output):
        seen_lm["done"] = True
    lm_hook_handles.append(model.llm.norm.register_forward_hook(flip_done))

    # Per-layer hooks above already cover layers 0-6, 13-20 and final norm.
    # Add the embed boundary on top and a sub-module bisect inside layer 1 to
    # locate where the L0 -> L1 jump (15x in one layer) originates.
    orig_prepare_embed = model._prepare_embed_inputs
    def hooked_prepare_embed(input_ids, audio_mask):
        out = orig_prepare_embed(input_ids, audio_mask)
        if not seen_lm["done"]:
            arr = out.detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, "lm-hidden-step0-cond-embed.bin"), arr[0])
            if arr.shape[0] >= 2:
                save_dump(os.path.join(dump_dir, "lm-hidden-step0-uncond-embed.bin"), arr[1])
        return out
    model._prepare_embed_inputs = hooked_prepare_embed

    # Layer 1 sub-module bisect : dump after input_layernorm, after self_attn
    # (pre-residual), after post_attn_layernorm and after mlp (pre-residual).
    # Suffixes : -l1-norm1, -l1-attn, -l1-norm2, -l1-mlp.
    def make_sub_hook(suffix):
        def hook(module, inputs, output):
            if seen_lm["done"]:
                return
            h = output[0] if isinstance(output, tuple) else output
            if not isinstance(h, torch.Tensor) or h.ndim < 3:
                return
            arr = h.detach().to(torch.float32).cpu().numpy()
            save_dump(os.path.join(dump_dir, f"lm-hidden-step0-cond-l1-{suffix}.bin"), arr[0])
            if arr.shape[0] >= 2:
                save_dump(os.path.join(dump_dir, f"lm-hidden-step0-uncond-l1-{suffix}.bin"), arr[1])
        return hook
    if len(model.llm.layers) > 1:
        l1 = model.llm.layers[1]
        lm_hook_handles.append(l1.input_layernorm.register_forward_hook(make_sub_hook("norm1")))
        lm_hook_handles.append(l1.self_attn.register_forward_hook(make_sub_hook("attn")))
        lm_hook_handles.append(l1.post_attention_layernorm.register_forward_hook(make_sub_hook("norm2")))
        lm_hook_handles.append(l1.mlp.register_forward_hook(make_sub_hook("mlp")))

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt",    default="../examples/prompt.txt")
    ap.add_argument("--ref-text",  default="../examples/freeman.txt")
    ap.add_argument("--ref-wav",   default="../examples/freeman.wav")
    ap.add_argument("--seed",      type=int, default=42)
    ap.add_argument("--lang",      default="English")
    ap.add_argument("--duration",  type=float, default=None)
    ap.add_argument("--quant",     default="F32",
                    help="quantization suffix for GGUF (default: F32, e.g. BF16, Q8_0, Q4_K_M)")
    ap.add_argument("--out-cpp",   default="cpp/clone-cpp.wav")
    ap.add_argument("--out-pt",    default="python/clone-python.wav")
    args = ap.parse_args()

    model_lm  = MODEL_LM_T.format(q=args.quant)
    model_cdc = MODEL_CDC_T.format(q=args.quant)
    for p in (model_lm, model_cdc):
        if not os.path.isfile(p):
            print(f"[Error] GGUF not found: {p}")
            sys.exit(1)
    print(f"[Quant] {args.quant} -> {model_lm} + {model_cdc}")

    ensure_dir(DUMP_CPP)
    ensure_dir(DUMP_PT)
    os.makedirs(os.path.dirname(args.out_cpp) or ".", exist_ok=True)

    with open(args.prompt, "r", encoding="utf-8") as f:
        text = f.read().strip()
    with open(args.ref_text, "r", encoding="utf-8") as f:
        ref_text = f.read().strip()
    print(f"[Input] Prompt: {len(text)} chars: {text[:60]}{'...' if len(text) > 60 else ''}")
    print(f"[Input] RefText: {len(ref_text)} chars: {ref_text[:60]}{'...' if len(ref_text) > 60 else ''}")
    print(f"[Input] RefWav: {args.ref_wav}")
    print(f"[Input] Language: {args.lang}")
    print(f"[Input] Seed: {args.seed}")

    fix_random_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model  = OmniVoice.from_pretrained(
        CKPT,
        torch_dtype=torch.float32,
        attn_implementation="eager",
    ).to(device).eval()
    install_hooks(model, DUMP_PT)
    gen_kwargs = dict(
        text=text,
        language=args.lang,
        ref_text=ref_text,
        ref_audio=args.ref_wav,
        duration=args.duration,
    )
    audios = model.generate(**gen_kwargs)
    audio_pt = np.asarray(audios[0], dtype=np.float32)
    sf.write(args.out_pt, audio_pt, 24000, subtype="FLOAT")
    print(f"[Python] Audio: {audio_pt.shape[0]} samples {audio_pt.shape[0] / 24000:.2f}s -> {args.out_pt}")

    del model
    torch.cuda.empty_cache()

    # Strict F32 conformance: disable FlashAttention (eager attention) and
    # force strict F32 cuBLAS to match the Python side with allow_tf32=False
    # set at the top of this file. Both pipelines encode the reference audio
    # independently so every encoder stage gets compared frontally.
    cmd = [
        BIN,
        "--model",         model_lm,
        "--codec",         model_cdc,
        "--seed",          str(args.seed),
        "--ref-wav",       args.ref_wav,
        "--ref-text",      args.ref_text,
        "--lang",          args.lang,
        "--format",        "wav32",
        "--dump",          DUMP_CPP,
        "-o",              args.out_cpp,
        "--no-fa",
    ]
    if args.duration:
        cmd += ["--duration", str(args.duration)]
    print(f"[GGML] Cmd: {' '.join(cmd)}")
    r = subprocess.run(cmd, input=text, text=True)
    if r.returncode != 0:
        sys.exit(r.returncode)
    audio_cpp, sr = sf.read(args.out_cpp)
    if audio_cpp.ndim > 1:
        audio_cpp = audio_cpp[:, 0]
    audio_cpp = audio_cpp.astype(np.float32)
    print(f"[GGML] Audio: {audio_cpp.shape[0]} samples {sr} Hz {audio_cpp.shape[0] / sr:.2f}s -> {args.out_cpp}")

    # Cossim in pipeline order: prompt-ids -> logits -> tokens -> audio.
    # Cond and uncond on the same line so a drift localizes to the originating
    # stage. PromptIDs split into zones (style, text, ref-audio, target) when
    # the lengths can be inferred from the prompt log.
    def pair(name):
        a, _ = load_dump(os.path.join(DUMP_CPP, name))
        b, _ = load_dump(os.path.join(DUMP_PT,  name))
        return a, b

    ca, cb = pair("prompt-cond-ids.bin")
    ua, ub = pair("prompt-uncond-ids.bin")
    n = min(ca.size, cb.size)
    cai = ca.astype(np.int64).ravel()[:n]
    cbi = cb.astype(np.int64).ravel()[:n]
    cdiffs = np.where(cai != cbi)[0]
    n2 = min(ua.size, ub.size)
    uai = ua.astype(np.int64).ravel()[:n2]
    ubi = ub.astype(np.int64).ravel()[:n2]
    udiffs = np.where(uai != ubi)[0]
    print(f"[Cossim] PromptIDs cond exact: {100.0 * (1 - cdiffs.size / max(n, 1)):.2f}% diffs: {cdiffs.size} uncond exact: {100.0 * (1 - udiffs.size / max(n2, 1)):.2f}% diffs: {udiffs.size}")
    if cdiffs.size > 0:
        print(f"[Cossim] PromptIDs cond first diff range: s={cdiffs[0]}..{cdiffs[-1]}")
    if udiffs.size > 0:
        print(f"[Cossim] PromptIDs uncond first diff range: s={udiffs[0]}..{udiffs[-1]}")

    def maxabs(a, b):
        a = a.astype(np.float64).ravel()
        b = b.astype(np.float64).ravel()
        n = min(len(a), len(b))
        if not n:
            return 0.0
        a, b = a[:n], b[:n]
        # log_probs has -inf at the audio_mask_id slot on both sides. Drop
        # positions that are non finite on either side so the subtract is safe.
        keep = np.isfinite(a) & np.isfinite(b)
        if not keep.all():
            a = a[keep]
            b = b[keep]
        return float(np.max(np.abs(a - b))) if a.size else 0.0

    ra, rb = pair("ref-audio-16k.bin")
    print(f"[Cossim] RefAudio16k max_abs_diff: {maxabs(ra, rb):.3e} cossim: {cos(ra, rb):.6f} samples: {min(ra.size, rb.size)}")

    fa, fb = pair("ref-hubert-features.bin")
    print(f"[Cossim] HuBERT-features cossim: {cos(fa, fb):.6f} max_abs_diff: {maxabs(fa, fb):.3e} shape_cpp: {fa.shape} shape_pt: {fb.shape}")

    # HuBERT bisect taps in pipeline order : feat-extract (post 7 conv1d) ->
    # feat-proj-ln (post LN, pre Linear) -> feat-proj (post 512 -> 768 Linear)
    # -> enc-init (post pos_conv add + LN) -> l0 -> l5 -> l7 -> l9 -> l11. The
    # downsample 2x and the 13-state mean that follow are already covered by
    # HuBERT-features above.
    for tap in ["hubert-feat-extract", "hubert-feat-proj-ln", "hubert-feat-proj",
                "hubert-enc-init", "hubert-l0", "hubert-l5", "hubert-l7", "hubert-l9", "hubert-l11"]:
        ta, tb = pair(tap + ".bin")
        print(f"[Cossim] {tap} cossim: {cos(ta, tb):.6f} max_abs_diff: {maxabs(ta, tb):.3e} shape_cpp: {ta.shape} shape_pt: {tb.shape}")

    ka, kb = pair("ref-audio-codes.bin")
    n_k = min(ka.size, kb.size)
    kai = ka.astype(np.int64).ravel()[:n_k]
    kbi = kb.astype(np.int64).ravel()[:n_k]
    kdiffs = np.where(kai != kbi)[0]
    print(f"[Cossim] RefAudioCodes exact: {100.0 * (1 - kdiffs.size / max(n_k, 1)):.2f}% diffs: {kdiffs.size} shape_cpp: {ka.shape} shape_pt: {kb.shape}")
    if ka.ndim == 2 and kb.ndim == 2 and ka.shape == kb.shape:
        per_k = []
        for k in range(ka.shape[0]):
            d = int(np.sum(ka[k].astype(np.int64) != kb[k].astype(np.int64)))
            per_k.append(f"k{k}={100.0 * (1 - d / ka.shape[1]):.1f}%")
        print(f"[Cossim] RefAudioCodes per-codebook: {' '.join(per_k)}")

    # Per layer LM hidden state drift bisect. C++ dumps cond and uncond at
    # 4 mid layers and the final norm. Pairs only print if both files exist.
    for layer_name, label in [("l0", "L0"),
                              ("l1-norm1", "L1.norm1"), ("l1-attn", "L1.attn"),
                              ("l1-norm2", "L1.norm2"), ("l1-mlp", "L1.mlp"),
                              ("l1", "L1"), ("l2", "L2"), ("l3", "L3"), ("l4", "L4"), ("l5", "L5"), ("l6", "L6"),
                              ("l13", "L13"), ("l14", "L14"), ("l15", "L15"), ("l16", "L16"), ("l17", "L17"), ("l18", "L18"),
                              ("l19", "L19"), ("l20", "L20"), ("", "Lf")]:
        suffix = f"-{layer_name}" if layer_name else ""
        cn = f"lm-hidden-step0-cond{suffix}.bin"
        un = f"lm-hidden-step0-uncond{suffix}.bin"
        try:
            ca, cb = pair(cn)
            ua, ub = pair(un)
        except FileNotFoundError:
            continue
        # Trim cpp shape (which is c_len) into the larger pt shape (max_c_len)
        nc = min(ca.size, cb.size)
        nu = min(ua.size, ub.size)
        cad = ca.ravel()[:nc].astype(np.float64)
        cbd = cb.ravel()[:nc].astype(np.float64)
        uad = ua.ravel()[:nu].astype(np.float64)
        ubd = ub.ravel()[:nu].astype(np.float64)
        max_c = float(np.max(np.abs(cad - cbd))) if nc else 0.0
        max_u = float(np.max(np.abs(uad - ubd))) if nu else 0.0
        cos_c = cos(cad, cbd)
        cos_u = cos(uad, ubd)
        print(f"[Cossim] {label} hidden cond: cos={cos_c:.6f} max_abs={max_c:.3e}  uncond: cos={cos_u:.6f} max_abs={max_u:.3e}")

    ca, cb = pair("lm-logits-step0-cond.bin")
    ua, ub = pair("lm-logits-step0-uncond.bin")
    print(f"[Cossim] Logits cond: {cos(ca, cb):.6f} uncond: {cos(ua, ub):.6f}")

    la, lb = pair("mg-log-probs-step0.bin")
    print(f"[Cossim] Step0 log_probs cossim: {cos(la, lb):.6f} max_abs_diff: {maxabs(la, lb):.3e}")

    pa, pb = pair("mg-pred-tokens-step0.bin")
    n_p = min(pa.size, pb.size)
    pai = pa.astype(np.int64).ravel()[:n_p]
    pbi = pb.astype(np.int64).ravel()[:n_p]
    pdiffs = np.where(pai != pbi)[0]
    print(f"[Cossim] Step0 pred_tokens exact: {100.0 * (1 - pdiffs.size / max(n_p, 1)):.2f}% diffs: {pdiffs.size}/{n_p}")

    sa, sb = pair("mg-scores-step0.bin")
    print(f"[Cossim] Step0 scores cossim: {cos(sa, sb):.6f} max_abs_diff: {maxabs(sa, sb):.3e}")

    ta, tb = pair("mg-tokens.bin")
    n = min(ta.size, tb.size)
    ai = ta.astype(np.int64).ravel()[:n]
    bi = tb.astype(np.int64).ravel()[:n]
    print(f"[Cossim] Tokens: {cos(ta, tb):.6f} exact: {100.0 * float(np.mean(ai == bi)):.2f}%")

    aa, ab = pair("output-audio.bin")
    print(f"[Cossim] Audio: {cos(aa, ab):.6f}")

    n = min(audio_cpp.size, audio_pt.size)
    print(f"[Cossim] WAV stft_cos: {stft_cos(audio_cpp[:n], audio_pt[:n]):.6f} samples: {n}")

if __name__ == "__main__":
    main()