#!/usr/bin/env python3
"""
Export Qwen3-TTS speech decoder V2 to full ONNX for GPU QNN backend.
Exact architecture from state dict analysis:
  pre_conv: Conv1d(512, 1024, kernel=3, pad=1)
  pre_transformer: 8-layer transformer (dim=512, attn_dim=1024, heads=16, head_dim=64)
  upsample: 2× ConvTranspose1d(1024,1024) + ConvNeXtV2
  BigVGAN: 7 blocks with Snake activation, upsample_rates=[8,5,4,3]
"""
import sys, os, math, warnings
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
warnings.filterwarnings("ignore")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

state = torch.load("/opt/Kazeia/models_qnn/qwen3-tts-native/speech_decoder_weights.pt", map_location="cpu")
print(f"Loaded {len(state)} tensors")

OUT = "/opt/Kazeia/models_qnn/qwen3-tts-decoder-full-onnx"
for d in ["pre_conv", "preprocessor", "conv_decoder"]:
    os.makedirs(f"{OUT}/{d}", exist_ok=True)

SEQ = 60

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

# ===== DiT Layer =====
class DiTLayer(nn.Module):
    def __init__(self, dim=512, attn_dim=1024, n_heads=16, head_dim=64, ffn_dim=1024, rope_theta=10000.0):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.input_layernorm = RMSNorm(dim)
        self.post_attention_layernorm = RMSNorm(dim)
        self.self_attn = nn.Module()
        self.self_attn.q_proj = nn.Linear(dim, attn_dim, bias=False)
        self.self_attn.k_proj = nn.Linear(dim, attn_dim, bias=False)
        self.self_attn.v_proj = nn.Linear(dim, attn_dim, bias=False)
        self.self_attn.o_proj = nn.Linear(attn_dim, dim, bias=False)
        self.self_attn_layer_scale = nn.Module()
        self.self_attn_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01)
        self.mlp = nn.Module()
        self.mlp.gate_proj = nn.Linear(dim, ffn_dim, bias=False)
        self.mlp.up_proj = nn.Linear(dim, ffn_dim, bias=False)
        self.mlp.down_proj = nn.Linear(ffn_dim, dim, bias=False)
        self.mlp_layer_scale = nn.Module()
        self.mlp_layer_scale.scale = nn.Parameter(torch.ones(dim) * 0.01)
        # Pre-compute RoPE
        freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(512).float()
        emb = torch.outer(t, freqs)
        self.register_buffer("rope_cos", emb.cos())
        self.register_buffer("rope_sin", emb.sin())

    def apply_rope(self, x, seq_len):
        cos = self.rope_cos[:seq_len].unsqueeze(0).unsqueeze(0)  # [1,1,T,D/2]
        sin = self.rope_sin[:seq_len].unsqueeze(0).unsqueeze(0)
        x1, x2 = x[..., ::2], x[..., 1::2]
        out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
        return out.flatten(-2)

    def forward(self, x):
        B, T, _ = x.shape
        h = self.input_layernorm(x)
        q = self.self_attn.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.self_attn.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.self_attn.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        q = self.apply_rope(q, T)
        k = self.apply_rope(k, T)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.full((T, T), float('-inf'), device=x.device), diagonal=1)
        attn = attn + mask.unsqueeze(0).unsqueeze(0)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(B, T, -1)
        x = x + self.self_attn_layer_scale.scale * self.self_attn.o_proj(out)
        h2 = self.post_attention_layernorm(x)
        ffn = self.mlp.down_proj(F.silu(self.mlp.gate_proj(h2)) * self.mlp.up_proj(h2))
        x = x + self.mlp_layer_scale.scale * ffn
        return x

# ===== Pre-transformer =====
class PreTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.input_proj = nn.Linear(1024, 512)
        self.layers = nn.ModuleList([DiTLayer() for _ in range(8)])
        self.norm = RMSNorm(512)
        self.output_proj = nn.Linear(512, 1024)
    def forward(self, x):
        x = x.transpose(1, 2)  # [B,C,T] → [B,T,C]
        x = self.input_proj(x)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        x = self.output_proj(x)
        return x.transpose(1, 2)  # [B,T,C] → [B,C,T]

# ===== ConvNeXt-V2 =====
class ConvNeXtV2(nn.Module):
    def __init__(self, dim=1024, int_dim=4096):
        super().__init__()
        self.dwconv = nn.Module()
        self.dwconv.conv = nn.Conv1d(dim, dim, 7, padding=3, groups=dim)
        self.norm = nn.LayerNorm(dim)
        self.pwconv1 = nn.Linear(dim, int_dim)
        self.pwconv2 = nn.Linear(int_dim, dim)
        self.gamma = nn.Parameter(torch.ones(dim) * 0.01)
    def forward(self, x):
        r = x
        x = self.dwconv.conv(x)
        x = x.transpose(1, 2)
        x = self.norm(x)
        x = F.gelu(self.pwconv1(x))
        x = self.pwconv2(x)
        x = (self.gamma * x).transpose(1, 2)
        return r + x

# ===== Upsample =====
class Upsample(nn.Module):
    def __init__(self):
        super().__init__()
        # upsampling_ratios=[2,2] from config
        # ConvTranspose1d(1024,1024,kernel=2*ratio,stride=ratio,padding=ratio//2)
        w0 = state["upsample.0.0.conv.weight"]
        w1 = state["upsample.1.0.conv.weight"]
        # For upsampling ratio r: kernel=2*r, stride=r, pad=r//2
        # But kernel=2 could mean ratio=1 or ratio=2 with pad adjustment
        # From config: upsampling_ratios=[2,2], so stride=2 for both
        # kernel=2, stride=2, padding=0 gives exact 2× upsample
        self.conv0 = nn.ConvTranspose1d(w0.shape[0], w0.shape[1], w0.shape[2], stride=2, padding=0)
        self.cnx0 = ConvNeXtV2()
        self.conv1 = nn.ConvTranspose1d(w1.shape[0], w1.shape[1], w1.shape[2], stride=2, padding=0)
        self.cnx1 = ConvNeXtV2()
        print(f"  Upsample 0: {w0.shape} stride=2")
        print(f"  Upsample 1: {w1.shape} stride=2")
    def forward(self, x):
        x = self.conv0(x)
        x = self.cnx0(x)
        x = self.conv1(x)
        x = self.cnx1(x)
        return x

# ===== SnakeBeta =====
class SnakeBeta(nn.Module):
    def __init__(self, channels):
        super().__init__()
        # Stored as [channels] in checkpoint, reshaped to [1,channels,1] in forward
        self.alpha = nn.Parameter(torch.ones(channels))
        self.beta = nn.Parameter(torch.ones(channels))
    def forward(self, x):
        a = self.alpha.view(1, -1, 1)
        b = self.beta.view(1, -1, 1)
        return x + (1.0 / (b + 1e-9)) * (torch.sin(a * x) ** 2)

# ===== AMPBlock (BigVGAN resblock) =====
class AMPBlock(nn.Module):
    def __init__(self, channels, kernel_size=7):
        super().__init__()
        self.act1 = SnakeBeta(channels)
        self.act2 = SnakeBeta(channels)
        self.conv1 = nn.Module()
        self.conv1.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size//2)
        self.conv2 = nn.Module()
        self.conv2.conv = nn.Conv1d(channels, channels, kernel_size, padding=1)
    def forward(self, x):
        h = self.act1(x)
        h = self.conv1.conv(h)
        h = self.act2(h)
        h = self.conv2.conv(h)
        return x + h

# ===== UpsampleBigVGANBlock =====
class UpsampleBigVGANBlock(nn.Module):
    def __init__(self, ch_in, ch_out, kernel, rate, pad, block_idx):
        super().__init__()
        self.block = nn.ModuleList()
        self.block.append(SnakeBeta(ch_in))
        self.block.append(nn.ConvTranspose1d(ch_in, ch_out, kernel, stride=rate, padding=pad))
        for j in range(3):
            amp_w1 = state[f"decoder.{block_idx}.block.{j+2}.conv1.conv.weight"]
            amp_w2 = state[f"decoder.{block_idx}.block.{j+2}.conv2.conv.weight"]
            amp = AMPBlock(amp_w1.shape[0], amp_w1.shape[2])
            amp.conv2.conv = nn.Conv1d(amp_w2.shape[1], amp_w2.shape[0], amp_w2.shape[2], padding=amp_w2.shape[2]//2)
            self.block.append(amp)
    def forward(self, x):
        # block[0] = SnakeBeta, block[1] = ConvTranspose1d, block[2-4] = AMPBlocks
        x = self.block[0](x)
        x = self.block[1](x)
        for i in range(2, 5):
            x = self.block[i](x)
        return x

# ===== BigVGAN =====
class BigVGAN(nn.Module):
    def __init__(self):
        super().__init__()
        # decoder.0: initial conv [1536, 1024, 7]
        w0 = state["decoder.0.conv.weight"]
        ch = w0.shape[0]  # 1536
        self.initial_conv = nn.Module()
        self.initial_conv.conv = nn.Conv1d(w0.shape[1], ch, w0.shape[2], padding=w0.shape[2]//2)

        # decoder.1-4: upsample blocks
        # Each: SnakeBeta(ch_in) + ConvTranspose1d(ch_in, ch_out) + 3×AMPBlock(ch_out)
        self.blocks = nn.ModuleList()
        for i in range(4):
            block_idx = i + 1
            conv_w = state[f"decoder.{block_idx}.block.1.conv.weight"]
            ch_in = conv_w.shape[0]   # input channels for this ConvTranspose1d (= ch from snake)
            ch_out = conv_w.shape[1]  # output channels (Note: ConvTranspose1d weight is [ch_in, ch_out, K])
            kernel = conv_w.shape[2]
            # Determine stride and padding from kernel and upsample rate
            # The upsample rates are [8,5,4,3], kernel = rate*2
            rate = [8, 5, 4, 3][i]
            pad = (kernel - rate) // 2

            blk = UpsampleBigVGANBlock(ch_in, ch_out, kernel, rate, pad, block_idx)
            self.blocks.append(blk)
            ch = ch_out

        # decoder.5: final SnakeBeta
        self.final_act = SnakeBeta(ch)
        # decoder.6: final conv to audio
        w6 = state["decoder.6.conv.weight"]
        self.final_conv = nn.Module()
        self.final_conv.conv = nn.Conv1d(w6.shape[1], w6.shape[0], w6.shape[2], padding=w6.shape[2]//2)

    def forward(self, x):
        x = self.initial_conv.conv(x)
        for blk in self.blocks:
            for m in blk.block:
                x = m(x)
        x = self.final_act(x)
        x = self.final_conv.conv(x)
        x = torch.tanh(x)
        return x

# ===== Build & Load =====
print("\n--- Building pre_conv ---")
pre_conv = nn.Conv1d(512, 1024, 3, padding=1)
pre_conv.weight.data = state["pre_conv.conv.weight"]
pre_conv.bias.data = state["pre_conv.conv.bias"]
pre_conv.eval()
dummy = torch.randn(1, 512, SEQ)
with torch.no_grad():
    pc = pre_conv(dummy)
print(f"pre_conv: [1,512,{SEQ}] → {list(pc.shape)}")

print("\n--- Building pre_transformer ---")
pt = PreTransformer()
pt_state = {k[len("pre_transformer."):]: v for k, v in state.items() if k.startswith("pre_transformer.")}
m, u = pt.load_state_dict(pt_state, strict=False)
print(f"pre_transformer load: missing={len(m)}, unexpected={len(u)}")
if m: print(f"  Missing: {m[:3]}")
pt.eval()
with torch.no_grad():
    pt_out = pt(pc)
print(f"pre_transformer: {list(pc.shape)} → {list(pt_out.shape)}")

print("\n--- Building upsample ---")
up = Upsample()
up_map = {}
for k, v in state.items():
    if k.startswith("upsample."):
        nk = k[len("upsample."):]
        # nk = "0.0.conv.weight" → conv0.weight, "0.1.dwconv.conv.weight" → cnx0.dwconv.conv.weight
        parts = nk.split(".", 2)
        idx = parts[0]  # 0 or 1
        sub = parts[1]  # 0 (conv) or 1 (convnext)
        rest = parts[2] if len(parts) > 2 else ""
        if sub == "0":
            # conv transpose: upsample.0.0.conv.weight → conv0.weight (remove ".conv" layer)
            if rest.startswith("conv."):
                up_map[f"conv{idx}.{rest[5:]}"] = v
            else:
                up_map[f"conv{idx}.{rest}"] = v
        else:
            up_map[f"cnx{idx}.{rest}"] = v
m, u = up.load_state_dict(up_map, strict=False)
print(f"upsample load: missing={len(m)}, unexpected={len(u)}")
if m: print(f"  Missing: {m[:3]}")
up.eval()
with torch.no_grad():
    up_out = up(pt_out)
print(f"upsample: {list(pt_out.shape)} → {list(up_out.shape)}")

print("\n--- Building BigVGAN ---")
bv = BigVGAN()
bv_map = {}
for k, v in state.items():
    if k.startswith("decoder."):
        nk = k[len("decoder."):]
        parts = nk.split(".", 1)
        idx = int(parts[0])
        rest = parts[1] if len(parts) > 1 else ""
        if idx == 0:
            # decoder.0.conv.weight → initial_conv.conv.weight
            bv_map[f"initial_conv.{rest}"] = v
        elif 1 <= idx <= 4:
            # decoder.{i}.block.1.conv.weight → blocks.{i-1}.block.1.weight (ConvTranspose1d, no ".conv")
            new_rest = rest
            if rest.startswith("block.1.conv."):
                # ConvTranspose1d: remove extra ".conv" layer
                new_rest = "block.1." + rest[len("block.1.conv."):]
            bv_map[f"blocks.{idx-1}.{new_rest}"] = v
        elif idx == 5:
            bv_map[f"final_act.{rest}"] = v
        elif idx == 6:
            bv_map[f"final_conv.{rest}"] = v
m, u = bv.load_state_dict(bv_map, strict=False)
print(f"BigVGAN load: missing={len(m)}, unexpected={len(u)}")
if m: print(f"  Missing: {m[:5]}")
if u: print(f"  Unexpected: {u[:5]}")
bv.eval()
with torch.no_grad():
    audio = bv(up_out)
print(f"BigVGAN: {list(up_out.shape)} → {list(audio.shape)}")

# ===== Export =====
print("\n=== Exporting pre_conv ===")
torch.onnx.export(pre_conv, dummy, f"{OUT}/pre_conv/model.onnx",
    input_names=["x"], output_names=["pre_conv_out"], opset_version=17)
print(f"  {os.path.getsize(f'{OUT}/pre_conv/model.onnx')/1024:.0f} KB")

class Preprocessor(nn.Module):
    def __init__(self, pt, up):
        super().__init__(); self.pt = pt; self.up = up
    def forward(self, x):
        return self.up(self.pt(x))

print("\n=== Exporting preprocessor ===")
prep = Preprocessor(pt, up).eval()
with torch.no_grad():
    prep_out = prep(pc)
torch.onnx.export(prep, pc, f"{OUT}/preprocessor/model.onnx",
    input_names=["pre_conv_out"], output_names=["hidden"], opset_version=17)
print(f"  {os.path.getsize(f'{OUT}/preprocessor/model.onnx')/1024/1024:.1f} MB")

print("\n=== Exporting conv_decoder ===")
torch.onnx.export(bv, up_out, f"{OUT}/conv_decoder/model.onnx",
    input_names=["hidden"], output_names=["audio"], opset_version=17)
print(f"  {os.path.getsize(f'{OUT}/conv_decoder/model.onnx')/1024/1024:.1f} MB")

# ===== Validate =====
print("\n=== Validation ===")
import onnxruntime as ort
for name, inp, ref in [("pre_conv", dummy, pc), ("preprocessor", pc, prep_out), ("conv_decoder", up_out, audio)]:
    sess = ort.InferenceSession(f"{OUT}/{name}/model.onnx")
    o = sess.run(None, {sess.get_inputs()[0].name: inp.detach().numpy()})[0]
    d = np.max(np.abs(o - ref.detach().numpy()))
    print(f"  {name}: max_diff={d:.6f} shape={o.shape}")

print("\nTotal:")
for n in ["pre_conv", "preprocessor", "conv_decoder"]:
    print(f"  {n}: {os.path.getsize(f'{OUT}/{n}/model.onnx')/1024/1024:.1f} MB")