#!/usr/bin/env python3
"""
Export Qwen3-TTS talker KV-cache to ONNX with M-RoPE.
Following TTS_CALIBRATION_GUIDE.md exactly:
  Inputs:  [emb(1,1,1024), mask(1,1,1,200), cos(1,1,128), sin(1,1,128), 56×kv(1,8,199,128)]
  Outputs: [hidden(1,1,1024), logits(1,1,3072), 56×kv(1,8,200,128)]
"""
import sys
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
import warnings; warnings.filterwarnings("ignore")
import torch, torch.nn as nn, torch.nn.functional as F
import numpy as np, os, math

N_LAYERS = 28; DIM = 1024; N_HEADS = 16; N_KV_HEADS = 8; HEAD_DIM = 128
KV_LEN = 199; VOCAB = 3072; FFN_DIM = 3072
GQA_GROUPS = N_HEADS // N_KV_HEADS  # 2

print("Loading talker weights...")
state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-export/qwen3_tts_talker.pth",
                    map_location="cpu", weights_only=False)

class RMSNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.weight

def apply_interleaved_rope(x, cos, sin):
    """Apply interleaved RoPE: x[..., ::2] and x[..., 1::2] rotated by cos/sin."""
    # x: [B, heads, 1, head_dim], cos/sin: [1, 1, 128] → use first 64 pairs
    half = x.shape[-1] // 2
    cos = cos[..., :half]  # [1,1,64]
    sin = sin[..., :half]
    x1 = x[..., ::2]   # [B, heads, 1, 64]
    x2 = x[..., 1::2]
    cos = cos.unsqueeze(1)  # [1, 1, 1, 64]
    sin = sin.unsqueeze(1)
    o1 = x1 * cos - x2 * sin
    o2 = x1 * sin + x2 * cos
    return torch.stack([o1, o2], dim=-1).flatten(-2)

class TalkerKVWrapper(nn.Module):
    def __init__(self, sd):
        super().__init__()
        self.norm = RMSNorm(DIM)
        self.norm.weight.data = sd["norm.weight"]
        self.lm_head = nn.Linear(DIM, VOCAB, bias=False)
        self.lm_head.weight.data = sd["output.weight"]

        self.layers = nn.ModuleList()
        for i in range(N_LAYERS):
            L = nn.Module()
            L.attn_norm = RMSNorm(DIM); L.attn_norm.weight.data = sd[f"layers.{i}.attention_norm.weight"]
            L.ffn_norm = RMSNorm(DIM); L.ffn_norm.weight.data = sd[f"layers.{i}.ffn_norm.weight"]
            L.wq = nn.Linear(DIM, N_HEADS*HEAD_DIM, bias=False); L.wq.weight.data = sd[f"layers.{i}.attention.wq.weight"]
            L.wk = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wk.weight.data = sd[f"layers.{i}.attention.wk.weight"]
            L.wv = nn.Linear(DIM, N_KV_HEADS*HEAD_DIM, bias=False); L.wv.weight.data = sd[f"layers.{i}.attention.wv.weight"]
            L.wo = nn.Linear(N_HEADS*HEAD_DIM, DIM, bias=False); L.wo.weight.data = sd[f"layers.{i}.attention.wo.weight"]
            L.q_norm = RMSNorm(HEAD_DIM); L.q_norm.weight.data = sd[f"layers.{i}.attention.q_norm_fn.weight"]
            L.k_norm = RMSNorm(HEAD_DIM); L.k_norm.weight.data = sd[f"layers.{i}.attention.k_norm_fn.weight"]
            L.w1 = nn.Linear(DIM, FFN_DIM, bias=False); L.w1.weight.data = sd[f"layers.{i}.feed_forward.w1.weight"]
            L.w2 = nn.Linear(FFN_DIM, DIM, bias=False); L.w2.weight.data = sd[f"layers.{i}.feed_forward.w2.weight"]
            L.w3 = nn.Linear(DIM, FFN_DIM, bias=False); L.w3.weight.data = sd[f"layers.{i}.feed_forward.w3.weight"]
            self.layers.append(L)

    def forward(self, emb, mask, cos, sin, *kv_in):
        """
        emb:  [1, 1, 1024]
        mask: [1, 1, 1, 200]  (right-aligned, -10000 for masked positions)
        cos:  [1, 1, 128]     (M-RoPE cos for current position)
        sin:  [1, 1, 128]     (M-RoPE sin for current position)
        kv_in: 28 × (k[1,8,199,128], v[1,8,199,128])
        Returns: hidden[1,1,1024], logits[1,1,3072], 28 × (k[1,8,200,128], v[1,8,200,128])
        """
        h = emb  # [1, 1, 1024]
        new_kvs = []

        for i, L in enumerate(self.layers):
            k_cache = kv_in[i * 2]      # [1, 8, KV_LEN, 128]
            v_cache = kv_in[i * 2 + 1]  # [1, 8, KV_LEN, 128]

            h_n = L.attn_norm(h)
            q = L.wq(h_n).view(1, 1, N_HEADS, HEAD_DIM).transpose(1, 2)    # [1,16,1,128]
            k = L.wk(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128]
            v = L.wv(h_n).view(1, 1, N_KV_HEADS, HEAD_DIM).transpose(1, 2) # [1,8,1,128]

            # q_norm, k_norm BEFORE rope
            q = L.q_norm(q)
            k = L.k_norm(k)

            # Apply M-RoPE (interleaved)
            q = apply_interleaved_rope(q, cos, sin)
            k = apply_interleaved_rope(k, cos, sin)

            # Concat new k,v to cache → [1, 8, 200, 128]
            k_full = torch.cat([k_cache, k], dim=2)  # [1, 8, KV_LEN+1, 128]
            v_full = torch.cat([v_cache, v], dim=2)

            # GQA expand for attention
            k_exp = k_full.repeat(1, GQA_GROUPS, 1, 1)  # [1, 16, 200, 128]
            v_exp = v_full.repeat(1, GQA_GROUPS, 1, 1)

            # Attention
            attn = torch.matmul(q, k_exp.transpose(-2, -1)) / math.sqrt(HEAD_DIM)  # [1,16,1,200]
            attn = attn + mask
            attn = F.softmax(attn, dim=-1)
            out = torch.matmul(attn, v_exp)  # [1,16,1,128]
            out = out.transpose(1, 2).contiguous().view(1, 1, N_HEADS * HEAD_DIM)  # [1,1,2048]
            h = h + L.wo(out)

            # FFN (SwiGLU)
            h_n2 = L.ffn_norm(h)
            h = h + L.w2(F.silu(L.w1(h_n2)) * L.w3(h_n2))

            new_kvs.extend([k_full, v_full])

        h = self.norm(h)
        logits = self.lm_head(h)  # [1, 1, 3072]

        return (h, logits, *new_kvs)

print("Building model...")
wrapper = TalkerKVWrapper(state)
wrapper.eval()

# Test with dummy inputs
print("Testing...")
dummy_emb = torch.randn(1, 1, DIM)
dummy_mask = torch.full((1, 1, 1, KV_LEN + 1), -10000.0)
dummy_mask[0, 0, 0, -1] = 0.0  # only current position visible
dummy_cos = torch.ones(1, 1, HEAD_DIM)
dummy_sin = torch.zeros(1, 1, HEAD_DIM)
kv_init = []
for _ in range(N_LAYERS):
    kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM))  # k
    kv_init.append(torch.zeros(1, N_KV_HEADS, KV_LEN, HEAD_DIM))  # v

with torch.no_grad():
    out = wrapper(dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init)

hidden = out[0]
logits = out[1]
print(f"hidden: {hidden.shape}, logits: {logits.shape}")
print(f"new_k[0]: {out[2].shape}, new_v[0]: {out[3].shape}")

# Export ONNX
OUT = "/opt/Kazeia/models_qnn/talker_kv_cpu_mrope"
os.makedirs(OUT, exist_ok=True)

input_names = ["inputs_embeds", "attention_mask", "cos", "sin"]
for i in range(N_LAYERS):
    input_names.extend([f"k_{i}_in", f"v_{i}_in"])

output_names = ["hidden", "logits"]
for i in range(N_LAYERS):
    output_names.extend([f"k_{i}_out", f"v_{i}_out"])

all_inputs = (dummy_emb, dummy_mask, dummy_cos, dummy_sin, *kv_init)

print(f"\nExporting ONNX ({len(input_names)} inputs, {len(output_names)} outputs)...")
torch.onnx.export(
    wrapper, all_inputs,
    f"{OUT}/model.onnx",
    input_names=input_names,
    output_names=output_names,
    opset_version=17,
)
sz = os.path.getsize(f"{OUT}/model.onnx")
print(f"Saved: {OUT}/model.onnx ({sz/1024/1024:.0f} MB)")

# Pre-compute M-RoPE cos/sin for all 200 positions and save
# For TTS, all 3 mrope axes use the same position, so M-RoPE = standard RoPE
# inv_freq = 1 / (theta^(2i/128)) for i in 0..63
print("\nPre-computing RoPE tables (standard, theta=1e6)...")
theta = 1000000.0
inv_freq = 1.0 / (theta ** (np.arange(0, HEAD_DIM, 2, dtype=np.float32) / HEAD_DIM))  # [64]
positions = np.arange(KV_LEN + 1, dtype=np.float32)  # [200]
angles = np.outer(positions, inv_freq)  # [200, 64]
rope_cos = np.cos(angles).astype(np.float32)  # [200, 64]
rope_sin = np.sin(angles).astype(np.float32)
# Expand to head_dim=128 for interleaved rope: each freq pair applies to (x[2i], x[2i+1])
rope_cos_full = np.repeat(rope_cos, 2, axis=1)  # [200, 128]
rope_sin_full = np.repeat(rope_sin, 2, axis=1)
np.save(f"{OUT}/talker_rotary_cos.npy", rope_cos_full)
np.save(f"{OUT}/talker_rotary_sin.npy", rope_sin_full)
print(f"Rotary tables: {rope_cos_full.shape} saved")

# Validate: run a quick prefill step and check output
print("\n=== Validation ===")
import onnxruntime as ort
sess = ort.InferenceSession(f"{OUT}/model.onnx")
print(f"ONNX inputs: {[i.name for i in sess.get_inputs()[:4]]}...")

mask = np.full((1, 1, 1, KV_LEN + 1), -10000.0, dtype=np.float32)
mask[0, 0, 0, -1] = 0.0
cos_0 = rope_cos_full[0:1].reshape(1, 1, HEAD_DIM)
sin_0 = rope_sin_full[0:1].reshape(1, 1, HEAD_DIM)

inputs = {
    "inputs_embeds": np.random.randn(1, 1, DIM).astype(np.float32),
    "attention_mask": mask,
    "cos": cos_0,
    "sin": sin_0,
}
for i in range(N_LAYERS):
    inputs[f"k_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32)
    inputs[f"v_{i}_in"] = np.zeros((1, N_KV_HEADS, KV_LEN, HEAD_DIM), dtype=np.float32)

out = sess.run(None, inputs)
print(f"ONNX hidden: {out[0].shape}, logits: {out[1].shape}")
print(f"ONNX k_0_out: {out[2].shape}")
print("Validation OK!")