#!/usr/bin/env python3
"""Extract embedding tables from Qwen3-TTS for the Android app."""
import torch
import numpy as np
import os

OUT = "/opt/Kazeia/tts_npy_exports"
os.makedirs(OUT, exist_ok=True)
NATIVE = "/opt/Kazeia/models_qnn/qwen3-tts-native"
TALKER_HF = "/opt/Kazeia/models_qnn/talker_hf"

# 1. Text embeddings + projection
print("=== text_components.pt ===")
tc = torch.load(f"{NATIVE}/text_components.pt", map_location="cpu", weights_only=False)
text_emb_dict = tc['text_embedding']  # OrderedDict
text_proj_dict = tc['text_projection']  # OrderedDict

print("text_embedding keys:", list(text_emb_dict.keys()))
print("text_projection keys:", list(text_proj_dict.keys()))

# text_embedding.weight: [151936, 2048] — full LLM vocab, but we only need first 1050 tokens
text_emb_w = text_emb_dict['weight']  # [151936, 2048]
print(f"text_embedding.weight: {text_emb_w.shape}")

# Only keep the first 1050 text tokens (CODEC_OFFSET in the engine)
text_emb_w = text_emb_w[:1050]  # [1050, 2048]
print(f"text_embedding (trimmed): {text_emb_w.shape}")

# text_projection is a 2-layer MLP: fc1(2048→1024) + GELU + fc2(1024→1024)
fc1_w = text_proj_dict['linear_fc1.weight']  # [hidden, 2048]
fc1_b = text_proj_dict['linear_fc1.bias']
fc2_w = text_proj_dict['linear_fc2.weight']  # [1024, hidden]
fc2_b = text_proj_dict['linear_fc2.bias']
print(f"text_projection: fc1={fc1_w.shape}, fc2={fc2_w.shape}")

# Apply: x = fc2(silu(fc1(text_emb)))  — hidden_act = "silu" from config.json
x = text_emb_w.float() @ fc1_w.float().T + fc1_b.float()
x = torch.nn.functional.silu(x)
text_projected = x @ fc2_w.float().T + fc2_b.float()
print(f"text_embeds_projected: {text_projected.shape}")
np.save(f"{OUT}/text_embeds_projected.npy", text_projected.numpy())

# 2. Codec embedding from talker HF safetensors
print("\n=== talker safetensors ===")
from safetensors.torch import load_file
talker_st = load_file(f"{TALKER_HF}/model.safetensors")
# Find the embed_tokens / codec_embedding
for k in sorted(talker_st.keys()):
    if 'embed' in k.lower():
        print(f"  {k}: {talker_st[k].shape}")

# The talker's token embedding is the codec embedding
codec_key = None
for k in talker_st.keys():
    if 'embed_tokens' in k or 'codec_embedding' in k:
        codec_key = k
        break
if codec_key:
    codec_emb = talker_st[codec_key].float().numpy()
    print(f"codec_embedding ({codec_key}): {codec_emb.shape}")
    np.save(f"{OUT}/codec_embedding.npy", codec_emb)
else:
    print("WARNING: codec_embedding not found in talker safetensors")
    # Fallback: try lm_head transposed
    for k in talker_st.keys():
        print(f"  {k}: {talker_st[k].shape}")

# 3. TTS special embeddings [3, 1024] = bos(2149), eos(2150), pad(2148)
# These are TEXT token IDs that go through text_embedding + text_projection (SiLU MLP)
# NOT codec_embedding! The IDs are in the LLM text vocab space (151936)
text_emb_full = text_emb_dict['weight']  # [151936, 2048] — full vocab
special_ids = [2149, 2150, 2148]  # bos, eos, pad
special_text = text_emb_full[special_ids].float()  # [3, 2048]
# Apply text_projection: fc1 + SiLU + fc2
h = special_text @ fc1_w.float().T + fc1_b.float()
h = torch.nn.functional.silu(h)
special_projected = (h @ fc2_w.float().T + fc2_b.float()).numpy()
np.save(f"{OUT}/tts_special_embeds.npy", special_projected)
print(f"tts_special_embeds: {special_projected.shape} (projected via text_projection SiLU)")

# 4. Code predictor embeddings — 15 codebook embedding tables [15, 2048, 1024]
print("\n=== code_predictor_weights.pt ===")
cp = torch.load(f"{NATIVE}/code_predictor_weights.pt", map_location="cpu", weights_only=False)
cp_embs = []
for i in range(15):
    key = f"model.codec_embedding.{i}.weight"
    if key in cp:
        cp_embs.append(cp[key].float().numpy())
if cp_embs:
    cp_all = np.stack(cp_embs)  # [15, 2048, 1024]
    np.save(f"{OUT}/code_predictor_embeddings.npy", cp_all)
    print(f"code_predictor_embeddings: {cp_all.shape}")

# 5. VQ codebooks — extract actual codebook vectors from embedding_sum / cluster_usage
print("\n=== speech_decoder_weights.pt (VQ codebooks) ===")
sd = torch.load(f"{NATIVE}/speech_decoder_weights.pt", map_location="cpu", weights_only=False)

# First codebook (rvq_first) + 15 rest codebooks (rvq_rest)
codebooks = []
# rvq_first
usage = sd['quantizer.rvq_first.vq.layers.0._codebook.cluster_usage']
emb_sum = sd['quantizer.rvq_first.vq.layers.0._codebook.embedding_sum']
# actual codebook = embedding_sum / cluster_usage (EMA-style)
cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy()
codebooks.append(cb)
print(f"  rvq_first codebook: {cb.shape}")

# rvq_rest has 15 layers (indices 0-14)
for i in range(15):
    usage = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.cluster_usage']
    emb_sum = sd[f'quantizer.rvq_rest.vq.layers.{i}._codebook.embedding_sum']
    cb = (emb_sum / usage.unsqueeze(1).clamp(min=1)).float().numpy()
    codebooks.append(cb)

vq_all = np.stack(codebooks)  # [16, 2048, 256]
np.save(f"{OUT}/vq_codebooks.npy", vq_all)
print(f"vq_codebooks: {vq_all.shape} (16 codebooks × 2048 entries × 256 dim)")

# Summary
print(f"\n=== Output: {OUT} ===")
for f in sorted(os.listdir(OUT)):
    sz = os.path.getsize(f"{OUT}/{f}")
    print(f"  {f}: {sz/1024/1024:.1f} MB")