#!/usr/bin/env python3
"""Cossim debug : C++ qwen-tts vs Python Qwen3-TTS on the Base 1.7B path.

Inputs (relative to CWD = tests/) :
    ../examples/prompt.txt       target text fed to both pipelines

Default mode is greedy (do_sample=False on both sides). The forward
chain is dumped layer by layer and compared paired with the Python
upstream hooks installed by cossim_common.install_hooks. Both pipelines
run on CUDA by default, the wrapper shell sweeps backends and quants.

Dumps land in cpp/base/ (C++) and python/base/ (Python). The script
compares each matching .bin pair via cosine similarity over the f32
payload, plus exact match rate for tensors that originated as int
(codec codes, prompt ids).
"""

import argparse
import os
import subprocess
import sys

import numpy as np
import soundfile as sf
import torch

import cossim_common as cc

MODEL_T     = "../models/qwen-talker-1.7b-base-{q}.gguf"
MODEL_CDC_T = "../models/qwen-tokenizer-12hz-{q}.gguf"
CKPT        = "../checkpoints/Qwen3-TTS-12Hz-1.7B-Base"
DUMP_CPP    = "cpp/base"
DUMP_PT     = "python/base"

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt",         default="../examples/prompt.txt")
    ap.add_argument("--seed",           type=int, default=42)
    ap.add_argument("--lang",           default="english")
    ap.add_argument("--quant",          default="F32",
                    help="GGUF quantization suffix (F32, BF16, Q8_0, Q4_K_M)")
    ap.add_argument("--out-pt",         default=os.path.join(DUMP_PT,  "base-python.wav"))
    ap.add_argument("--out-cpp",        default=os.path.join(DUMP_CPP, "base-cpp.wav"))
    ap.add_argument("--max-new-tokens", type=int, default=64)
    ap.add_argument("--trace",          action="store_true",
                    help="print per sample u and idx for the first 32 samples")
    args = ap.parse_args()

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

    with open(args.prompt, "r", encoding="utf-8") as f:
        text = f.read().strip()
    print(f"[Input] Prompt: {len(text)} chars: {text[:60]}{'...' if len(text) > 60 else ''}")
    print(f"[Input] Lang: {args.lang} Seed: {args.seed} MaxNewTokens: {args.max_new_tokens}")
    print(f"[Input] Mode: greedy")

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    cc.set_trace(args.trace)

    cc.register_qwen3_tts()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[Python] Device: {device}")
    model = cc.AutoModel.from_pretrained(
        CKPT,
        device_map=device,
        dtype=torch.float32,
        attn_implementation="eager",
    ).eval()
    processor = cc.AutoProcessor.from_pretrained(CKPT, fix_mistral_regex=True)

    assistant_text = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
    inp = processor(text=assistant_text, return_tensors="pt", padding=True)
    input_ids = inp["input_ids"].to(device)
    if input_ids.dim() == 1:
        input_ids = input_ids.unsqueeze(0)
    print(f"[Python] InputIds shape: {tuple(input_ids.shape)}")
    cc.save_dump_i32(os.path.join(DUMP_PT, "prompt-ids.bin"), input_ids[0])

    cc.install_hooks(model, DUMP_PT)

    # Custom subtalker_* kwargs are forwarded to talker.forward but not
    # declared on GenerationMixin, so transformers 4.57 rejects them under
    # the strict validator. Disable it on the talker only.
    model.talker._validate_model_kwargs = lambda *a, **k: None

    # Greedy hardcoded : argmax on both talker c0 and code predictor sub
    # codes. Stochastic mode is not exercised here because the F32 drift
    # between torch CUDA cuBLAS and ggml CUDA matmul on Qwen3 norm_w
    # inflated activations propagates through the FFN and flips multinomial
    # picks in flat distributions, breaking bit exactness. Argmax is robust
    # to that drift, so greedy gives 100 percent CodesFull match and
    # validates the full forward + sampling chain.
    gen_kwargs = dict(
        do_sample             = False,
        top_k                 = 1,
        top_p                 = 1.0,
        temperature           = 1.0,
        subtalker_dosample    = False,
        subtalker_top_k       = 1,
        subtalker_top_p       = 1.0,
        subtalker_temperature = 1.0,
        repetition_penalty    = 1.0,
    )

    talker_codes_list, _ = model.generate(
        input_ids=[input_ids],
        languages=[args.lang],
        non_streaming_mode=True,
        max_new_tokens=args.max_new_tokens,
        **gen_kwargs,
    )
    codes = talker_codes_list[0]
    print(f"[Python] Codes shape: {tuple(codes.shape)} (T_frames, num_code_groups)")
    cc.save_dump_i32(os.path.join(DUMP_PT, "codes-full.bin"),  codes)
    cc.save_dump_i32(os.path.join(DUMP_PT, "codes-step0.bin"), codes[0])

    wavs, fs = model.speech_tokenizer.decode([{"audio_codes": codes}])
    audio_pt = np.asarray(wavs[0], dtype=np.float32)
    sf.write(args.out_pt, audio_pt, fs, subtype="FLOAT")
    cc.save_dump(os.path.join(DUMP_PT, "output-audio.bin"), audio_pt)
    print(f"[Python] Audio: {audio_pt.shape[0]} samples {fs} Hz {audio_pt.shape[0]/fs:.2f}s -> {args.out_pt}")

    if not os.path.isfile(cc.BIN):
        print(f"[Cossim] FATAL: {cc.BIN} not found, build qwen-tts first")
        sys.exit(1)
    model_lm  = MODEL_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"[Cossim] FATAL: GGUF not found: {p}")
            sys.exit(1)
    print(f"[Quant] {args.quant} -> {model_lm} + {model_cdc}")

    del model
    if device == "cuda":
        torch.cuda.empty_cache()

    cmd = [
        cc.BIN,
        "--model",   model_lm,
        "--codec",   model_cdc,
        "--seed",    str(args.seed),
        "--text",    text,
        "--lang",    args.lang,
        "--max-new", str(args.max_new_tokens),
        "--dump",    DUMP_CPP,
        "-o",        args.out_cpp,
        "--greedy",
    ]
    print(f"[GGML] Cmd: {' '.join(cmd)}")
    r = subprocess.run(cmd)
    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}")

    cc.compare_exact_i32("prompt-ids.bin", DUMP_CPP, DUMP_PT, "PromptIDs")
    cc.compare_stages(cc.STAGES_STANDARD, DUMP_CPP, DUMP_PT)
    cc.compare_exact_i32("codes-full.bin", DUMP_CPP, DUMP_PT, "CodesFull")

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

    # STFT runs on the f32 bin dumps, not the WAV files : the C++ side writes
    # PCM_16 which quantizes the very low amplitudes of greedy short outputs
    # to zero, while the bin dumps preserve the raw float buffer.
    n = min(aa.size, ab.size)
    print(f"[Cossim] WAV stft_cos: {cc.stft_cos(aa.ravel()[:n], ab.ravel()[:n]):.6f} samples: {n}")

if __name__ == "__main__":
    main()