#!/usr/bin/env python3
# RLX — versatile ML compiler + runtime.
# Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

"""Dump talker layer-0 last-token intermediates (HF reference for eager parity)."""

from __future__ import annotations

import argparse
import json
import sys

import torch
from qwen_tts import Qwen3TTSModel


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model-dir", required=True)
    p.add_argument("--out", required=True)
    args = p.parse_args()

    m = Qwen3TTSModel.from_pretrained(args.model_dir, torch_dtype=torch.float32, device_map="cpu")
    text = "Hi."
    ids = m._tokenize_texts([m._build_assistant_text(text)])[0]

    talker_dec = m.model.talker.model
    embeds_holder: list[torch.Tensor] = []

    def pre_hook(_mod, _args, kwargs):
        ie = kwargs.get("inputs_embeds")
        if ie is not None and ie.shape[1] > 1:
            embeds_holder.append(ie.detach().clone())

    h = talker_dec.register_forward_pre_hook(pre_hook, with_kwargs=True)
    gen = m._merge_generate_kwargs(do_sample=False, subtalker_dosample=False, max_new_tokens=2)
    m.model.generate(
        input_ids=[ids],
        languages=["english"],
        speakers=["vivian"],
        non_streaming_mode=True,
        **gen,
    )
    h.remove()
    if not embeds_holder:
        print("no prefill embeds captured", file=sys.stderr)
        sys.exit(1)

    ie = embeds_holder[0]
    seq = ie.shape[1]
    attn = torch.ones(1, seq, dtype=torch.long)
    layer0 = talker_dec.layers[0]
    out: dict[str, list[float]] = {}

    def capture(name: str):
        def _hook(_mod, _inp, o):
            hs = o[0] if isinstance(o, tuple) else o
            out[name] = hs[0, -1].detach().cpu().tolist()

        return _hook

    hooks = [
        layer0.self_attn.register_forward_hook(capture("attn_out")),
        layer0.register_forward_hook(capture("layer_out")),
    ]
    with torch.no_grad():
        pos = torch.arange(seq, device=ie.device).view(1, 1, -1).expand(3, 1, -1).float()
        pe = talker_dec.rotary_emb(ie, pos)
        residual = ie
        h = layer0.input_layernorm(ie)
        attn_out, _ = layer0.self_attn(
            h, position_embeddings=pe, attention_mask=None, past_key_values=None
        )
        after_attn = (residual + attn_out)[0, -1]
        h2 = layer0.post_attention_layernorm(residual + attn_out)
        ff = layer0.mlp(h2)
        after_block = (residual + attn_out + ff)[0, -1]
        out["after_attn"] = after_attn.detach().cpu().tolist()
        out["after_block"] = after_block.detach().cpu().tolist()

    for hk in hooks:
        hk.remove()

    with open(args.out, "w", encoding="utf-8") as f:
        json.dump({"seq": seq, **out}, f)
    print(f"wrote {args.out}")


if __name__ == "__main__":
    main()
