"""Shared helpers for the qwentts.cpp cossim debug scripts.

Provides Philox uniform stream, dump load and save, install_hooks for the
talker submodel, the standard stage list and the metric helpers used by
debug-base / debug-tts / debug-customvoice / debug-clone cossim scripts.

Importing this module patches sys.path so qwen_tts upstream loads without
the V1 25Hz tokenizer (sox dependency stubbed out), and forces TF32 off on
every torch CUDA matmul path so Python results stay bit comparable across
runs and across machines.
"""

import os
import struct
import sys
import types

os.environ["NVIDIA_TF32_OVERRIDE"] = "0"

import numpy as np
import torch

torch.backends.cuda.matmul.allow_tf32                             = False
torch.backends.cudnn.allow_tf32                                   = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.set_float32_matmul_precision("highest")

# Stub the V1 25Hz tokenizer so qwen_tts.core imports without sox.
UPSTREAM_ROOT = "../../Qwen3-TTS"
sys.path.insert(0, UPSTREAM_ROOT)

class _StubV1Config:
    pass

class _StubV1Model:
    pass

_qwen_pkg          = types.ModuleType("qwen_tts")
_qwen_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts")]
sys.modules["qwen_tts"] = _qwen_pkg

_core_pkg          = types.ModuleType("qwen_tts.core")
_core_pkg.__path__ = [os.path.join(UPSTREAM_ROOT, "qwen_tts", "core")]
# Inject the stubbed core module before any submodule import so the real
# qwen_tts/core/__init__.py never runs : it pulls the V1 25Hz tokenizer that
# imports whisper_encoder, which prints a flash-attn warning at module load.
sys.modules["qwen_tts.core"] = _core_pkg
from qwen_tts.core.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2    import Qwen3TTSTokenizerV2Model
_core_pkg.Qwen3TTSTokenizerV1Config = _StubV1Config
_core_pkg.Qwen3TTSTokenizerV1Model  = _StubV1Model
_core_pkg.Qwen3TTSTokenizerV2Config = Qwen3TTSTokenizerV2Config
_core_pkg.Qwen3TTSTokenizerV2Model  = Qwen3TTSTokenizerV2Model

from qwen_tts.core.models.modeling_qwen3_tts      import Qwen3TTSForConditionalGeneration
from qwen_tts.core.models.configuration_qwen3_tts import Qwen3TTSConfig
from qwen_tts.core.models.processing_qwen3_tts    import Qwen3TTSProcessor
from transformers import AutoConfig, AutoModel, AutoProcessor
from transformers.utils import logging as hf_logging

# Silence the GenerationConfig validator that warns "flags are not valid
# and may be ignored" for temperature / top_k / top_p when do_sample=False.
# Those flags ride along inside the checkpoint generation_config and the
# greedy path drops them on purpose, the warning is just noise here.
hf_logging.set_verbosity_error()

# Register the Qwen3-TTS classes once per process. Calling twice raises a
# ValueError inside transformers, hence the guard.
_REGISTERED = {"done": False}

def register_qwen3_tts():
    if _REGISTERED["done"]:
        return
    AutoConfig.register("qwen3_tts", Qwen3TTSConfig)
    AutoModel.register(Qwen3TTSConfig, Qwen3TTSForConditionalGeneration)
    AutoProcessor.register(Qwen3TTSConfig, Qwen3TTSProcessor)
    _REGISTERED["done"] = True

# Path to the C++ qwen-tts binary, relative to tests/.
BIN = "../build/qwen-tts"

# Standard stage list shared by every cossim script. Mode-specific scripts
# may extend this list before iterating (eg. clone adds SpeakerEmb / RefCodes).
STAGES_STANDARD = [
    ("Embed",              "talker-input-embed.bin"),
    ("TrailingText",       "trailing-text-hidden.bin"),
    ("TTSPadEmbed",        "tts-pad-embed.bin"),
    ("L0",                 "talker-hidden-prefill-l0.bin"),
    ("L7",                 "talker-hidden-prefill-l7.bin"),
    ("L14",                "talker-hidden-prefill-l14.bin"),
    ("L21",                "talker-hidden-prefill-l21.bin"),
    ("L27",                "talker-hidden-prefill-l27.bin"),
    ("Final",              "talker-hidden-prefill-final.bin"),
    ("Logits",             "talker-logits-prefill.bin"),
    ("NextEmbStep0",       "next-emb-step0.bin"),
    ("TalkerHiddenStep1",  "talker-hidden-step1.bin"),
]

# Philox4x32-10 mirror of src/philox.h. Returns the same float u that
# philox_uniform_fill(seed, subseq, ctr_lo=0) gives for n=1.
PHILOX_M0     = np.uint32(0xD2511F53)
PHILOX_M1     = np.uint32(0xCD9E8D57)
PHILOX_W0     = np.uint32(0x9E3779B9)
PHILOX_W1     = np.uint32(0xBB67AE85)
TWO_POW32_INV = np.float32(2.3283064365386963e-10)

def _mulhilo32(a, b):
    p  = np.uint64(a) * np.uint64(b)
    lo = np.uint32(p & np.uint64(0xFFFFFFFF))
    hi = np.uint32(p >> np.uint64(32))
    return hi, lo

def _philox_round(ctr, k0, k1):
    hi0, lo0 = _mulhilo32(PHILOX_M0, ctr[0])
    hi1, lo1 = _mulhilo32(PHILOX_M1, ctr[2])
    return (np.uint32(hi1 ^ ctr[1] ^ k0),
            np.uint32(lo1),
            np.uint32(hi0 ^ ctr[3] ^ k1),
            np.uint32(lo0))

def _philox4x32_10(ctr, k0, k1):
    mask = np.uint64(0xFFFFFFFF)
    for _ in range(9):
        ctr = _philox_round(ctr, k0, k1)
        k0  = np.uint32((np.uint64(k0) + np.uint64(PHILOX_W0)) & mask)
        k1  = np.uint32((np.uint64(k1) + np.uint64(PHILOX_W1)) & mask)
    ctr = _philox_round(ctr, k0, k1)
    return ctr

def philox_uniform(seed, subseq, ctr_lo=0):
    slo = np.uint32(np.uint64(seed) & np.uint64(0xFFFFFFFF))
    shi = np.uint32(np.uint64(seed) >> np.uint64(32))
    ctr = (np.uint32(ctr_lo),
           np.uint32(0),
           np.uint32(np.uint64(subseq) & np.uint64(0xFFFFFFFF)),
           np.uint32(np.uint64(subseq) >> np.uint64(32)))
    r = _philox4x32_10(ctr, slo, shi)
    return (np.float32(r[0]) + np.float32(0.5)) * TWO_POW32_INV

# Globals advanced exactly once per multinomial sample, mirroring the C++
# side which advances subseq_counter at every sample_top_k_p call.
_subseq_counter = [0]
_seed           = [42]
_trace_samples  = [False]

def reset_philox(seed):
    _subseq_counter[0] = 0
    _seed[0]           = int(seed)

def set_trace(flag):
    _trace_samples[0] = bool(flag)

def patched_multinomial(input, num_samples, replacement=False, generator=None, out=None):
    """Drop in replacement for torch.multinomial(num_samples=1) that pulls
    the uniform draw from our Philox stream and walks the F32 cumulative
    sum the same way src/sampling.h does."""
    assert num_samples == 1, "patched_multinomial only handles num_samples=1"
    probs = input
    if probs.dim() == 1:
        probs = probs.unsqueeze(0)
    bsz, vocab = probs.shape
    out_ids = torch.zeros((bsz, 1), dtype=torch.long, device=probs.device)
    for b in range(bsz):
        u   = philox_uniform(_seed[0], _subseq_counter[0], 0)
        seq = _subseq_counter[0]
        _subseq_counter[0] += 1
        row = probs[b].to(torch.float32).cpu().numpy()
        s   = float(row.sum())
        # The C++ sampler draws u in [0, 1) and compares against acc/sum
        # implicitly via acc >= u*sum. We replicate that exact arithmetic.
        target = float(u) * s
        acc    = 0.0
        idx    = vocab - 1
        for i in range(vocab):
            acc += float(row[i])
            if acc >= target:
                idx = i
                break
        out_ids[b, 0] = idx
        if _trace_samples[0] and seq < 32:
            print(f"[Sample-PY] subseq={seq} u={float(u):.10f} idx={idx} top_prob={float(row.max()):.6f}")
    if input.dim() == 1:
        return out_ids.squeeze(0)
    return out_ids

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 save_dump_i32(path, data):
    if isinstance(data, torch.Tensor):
        data = data.detach().to(torch.int64).cpu().numpy()
    data  = np.ascontiguousarray(data.astype(np.int64))
    shape = data.shape
    fdata = data.astype(np.float32)
    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(fdata.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]
    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):
    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, bisect_layers=(0, 7, 14, 21, 27)):
    """Capture every intermediate tensor we can pin against the C++ side.
    Layout : bisection layers, final norm, prefill logits, the input embed
    fed to the talker and the trailing-text overlay tensors that drive the
    next-token embedding sums during generation. Counters track how many
    times the talker submodel has run so step 1 (first single-token forward
    after prefill) gets its own dump."""
    seen_prefill = {"done": False}
    seen_codes   = {"done": False}
    # talker_step counts how many times talker_model.forward has been called
    # after the prefill. Prefill itself is recorded as 0, the first single
    # token forward is 1, and so on. Single token forwards are detected by
    # inputs_embeds.shape[1] == 1 in the pre hook.
    talker_step  = {"n": 0}

    talker_model = model.talker.model
    talker_lm    = model.talker

    seen_layers = {idx: False for idx in bisect_layers}
    def make_layer_hook(layer_idx):
        def hook(module, inputs, output):
            if seen_layers[layer_idx]:
                return
            h = output[0] if isinstance(output, tuple) else output
            if h.dim() == 3:
                save_dump(os.path.join(dump_dir, f"talker-hidden-prefill-l{layer_idx}.bin"), h[0])
                seen_layers[layer_idx] = True
        return hook
    for layer_idx in bisect_layers:
        talker_model.layers[layer_idx].register_forward_hook(make_layer_hook(layer_idx))

    seen_norm = {"done": False}
    def norm_hook(module, inputs, output):
        if seen_norm["done"]:
            return
        if output.dim() == 3 and output.shape[1] > 1:
            save_dump(os.path.join(dump_dir, "talker-hidden-prefill-final.bin"), output[0])
            seen_norm["done"] = True
    talker_model.norm.register_forward_hook(norm_hook)

    # Pre-hook on the talker submodel : sees inputs_embeds whether the outer
    # talker forward was invoked with input_ids (single token step) or
    # inputs_embeds (prefill). The submodel always receives inputs_embeds
    # because the wrapper rebuilds it before calling self.model.
    def talker_model_pre_hook(module, args, kwargs):
        ie = kwargs.get("inputs_embeds", None)
        if ie is None:
            return
        if ie.dim() != 3:
            return
        if ie.shape[1] > 1:
            return
        if talker_step["n"] == 0:
            save_dump(os.path.join(dump_dir, "next-emb-step0.bin"), ie[0, 0])
        talker_step["n"] += 1
    talker_model.register_forward_pre_hook(talker_model_pre_hook, with_kwargs=True)

    # Post-hook on the talker submodel : captures last_hidden_state at step
    # 1 (first single token forward). That tensor is what feeds the code
    # predictor at step 1, so any drift between Python and C++ tells us the
    # next-emb-step0 changed the talker forward result.
    talker_post_step = {"n": 0}
    def talker_model_post_hook(module, inputs, output):
        last = output.last_hidden_state if hasattr(output, "last_hidden_state") else output[0]
        if last.dim() != 3 or last.shape[1] != 1:
            return
        if talker_post_step["n"] == 0:
            save_dump(os.path.join(dump_dir, "talker-hidden-step1.bin"), last[0, -1])
        talker_post_step["n"] += 1
    talker_model.register_forward_hook(talker_model_post_hook)

    # Talker LM wrapper hook : captures the prefill input embed (the talker
    # codec_embedding sum + text projection that mirrors what
    # prompt_builder_build produces in C++), the prefill logits, and the
    # trailing_text_hidden / tts_pad_embed overlay tensors carried by the
    # output dataclass at every step (we only dump them once).
    seen_overlay = {"done": False}
    orig_talker_forward = talker_lm.forward
    def hooked_talker_forward(*args, **kwargs):
        inputs_embeds = kwargs.get("inputs_embeds", None)
        if (inputs_embeds is not None and inputs_embeds.dim() == 3
                and inputs_embeds.shape[1] > 1 and not seen_prefill["done"]):
            save_dump(os.path.join(dump_dir, "talker-input-embed.bin"), inputs_embeds[0])
            seen_prefill["done"] = True
        out = orig_talker_forward(*args, **kwargs)
        if (out is not None and getattr(out, "logits", None) is not None
                and not seen_codes["done"]):
            logits = out.logits
            if logits.dim() == 3 and logits.shape[1] > 1:
                save_dump(os.path.join(dump_dir, "talker-logits-prefill.bin"), logits[0, -1])
                seen_codes["done"] = True
        if (out is not None and not seen_overlay["done"]
                and getattr(out, "trailing_text_hidden", None) is not None
                and getattr(out, "tts_pad_embed", None) is not None):
            tth = out.trailing_text_hidden
            tpe = out.tts_pad_embed
            if tth.dim() == 3 and tpe.dim() == 3:
                save_dump(os.path.join(dump_dir, "trailing-text-hidden.bin"), tth[0])
                save_dump(os.path.join(dump_dir, "tts-pad-embed.bin"), tpe[0, 0])
                seen_overlay["done"] = True
        return out
    talker_lm.forward = hooked_talker_forward

def pair(name, dump_cpp, dump_pt):
    a, _ = load_dump(os.path.join(dump_cpp, name))
    b, _ = load_dump(os.path.join(dump_pt,  name))
    return a, b

def metric(a, b):
    n     = min(a.size, b.size)
    af    = a.astype(np.float64).ravel()[:n]
    bf    = b.astype(np.float64).ravel()[:n]
    d     = np.abs(af - bf)
    nrm_a = float(np.linalg.norm(af))
    nrm_b = float(np.linalg.norm(bf))
    c     = float(np.dot(af, bf) / (nrm_a * nrm_b)) if nrm_a > 1e-10 and nrm_b > 1e-10 else 0.0
    return c, float(d.max()), float(d.mean())

def compare_stages(stages, dump_cpp, dump_pt):
    """Iterate the stages list and print one line per pair. Skips silently
    when a dump file is missing (eg. a mode that does not produce a given
    intermediate)."""
    for label, name in stages:
        try:
            a, b = pair(name, dump_cpp, dump_pt)
        except FileNotFoundError:
            print(f"[Cossim] {label} skipped (missing dump)")
            continue
        c, mx, mn = metric(a, b)
        print(f"[Cossim] {label} cos: {c:.6f} max: {mx:.4e} mean: {mn:.4e}")

def compare_exact_i32(name, dump_cpp, dump_pt, label):
    """Compare two int dumps stored as f32 (the encoding path used by both
    save_dump_i32 in Python and debug_dump_i32_as_f32 in C++). Prints an
    exact match percentage. Returns the percentage as a float."""
    a, b = pair(name, dump_cpp, dump_pt)
    ai   = a.astype(np.int64).ravel()
    bi   = b.astype(np.int64).ravel()
    n    = min(ai.size, bi.size)
    pct  = 100.0 * float(np.mean(ai[:n] == bi[:n]))
    print(f"[Cossim] {label} exact: {pct:.2f}% ({n} values)")
    return pct

# Greedy generation kwargs shared by every cossim script. do_sample=False
# alone selects argmax, top_k / top_p / temperature are intentionally
# omitted because GenerationConfig flags them as "not valid" warnings when
# do_sample=False. The subtalker_* keys are custom kwargs forwarded to the
# talker forward, the talker validator is bypassed by the script main.
GEN_KWARGS_GREEDY = dict(
    do_sample             = False,
    subtalker_dosample    = False,
    repetition_penalty    = 1.0,
)