#!/usr/bin/env python3
"""
export_cond_enc.py — Export AceStep Condition Encoder to ONNX.

The condition encoder takes outputs from the text encoder (text_hidden + lyric_embed)
and reference audio features (timbre_feats), and produces enc_hidden for DiT cross-attention.

Internal architecture:
  - text_projector: Linear(1024→2048, no bias) — projects text encoder output
  - lyric_encoder: Linear(1024→2048)+bias → 8-layer bidirectional Qwen3 → RMSNorm
  - timbre_encoder: Linear(64→2048)+bias → [CLS prepend] → 4-layer bidir Qwen3 → RMSNorm → position[0]
  - cat(lyric_out, timbre_out, text_proj_out) → enc_hidden [B, S_total, 2048]

Usage:
    python export_cond_enc.py --model-dir <path-to-DiT-safetensors> --output <output.onnx>
"""

import argparse
import os
import sys
import time
import struct
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn


class CondEncoderWrapper(nn.Module):
    """Wrapper for ONNX export that simplifies the condition encoder interface.
    
    For inference (batch_size=1), we simplify:
    - No pack_sequences sorting (all tokens are valid, no padding)
    - Timbre: single reference, so unpack is trivial (just unsqueeze)
    - Output is simple cat(lyric, timbre, text_proj)
    
    ONNX inputs:
        text_hidden:  [B, S_text, 1024]  fp16 — from text encoder
        lyric_embed:  [B, S_lyric, 1024] fp16 — from embedding table lookup
        timbre_feats: [B, S_ref, 64]     fp16 — from VAE encoder (or zeros)
        has_timbre:   [1]                int64 — 1 if timbre is present, 0 if not
    
    ONNX output:
        enc_hidden:   [B, S_total, 2048] fp16 — packed conditioning
    """
    
    def __init__(self, cond_encoder):
        super().__init__()
        self.text_projector = cond_encoder.text_projector
        self.lyric_encoder = cond_encoder.lyric_encoder
        self.timbre_encoder = cond_encoder.timbre_encoder
    
    def forward(self, text_hidden, lyric_embed, timbre_feats, has_timbre):
        """
        Forward pass with simplified interface for ONNX export.
        
        Note: For ONNX tracing, has_timbre must be a tensor, not a Python bool.
        We use torch.where / masking to handle the conditional timbre path.
        """
        B = text_hidden.shape[0]
        
        # 1) Text projection: [B, S_text, 1024] → [B, S_text, 2048]
        text_proj = self.text_projector(text_hidden)
        
        # 2) Lyric encoding: [B, S_lyric, 1024] → 8L bidir Qwen3 → [B, S_lyric, 2048]
        S_lyric = lyric_embed.shape[1]
        lyric_mask = torch.ones(B, S_lyric, device=lyric_embed.device, dtype=torch.long)
        lyric_out = self.lyric_encoder(
            inputs_embeds=lyric_embed,
            attention_mask=lyric_mask,
        )
        if hasattr(lyric_out, 'last_hidden_state'):
            lyric_out = lyric_out.last_hidden_state
        else:
            lyric_out = lyric_out[0]
        
        # 3) Timbre encoding: [B, S_ref, 64] → 4L bidir Qwen3 → position[0] → [B, 1, 2048]
        # For ONNX: we always run the timbre path but zero out if has_timbre=0
        S_ref = timbre_feats.shape[1]
        timbre_mask = torch.ones(1, S_ref, device=timbre_feats.device, dtype=torch.long)
        # refer_audio_order_mask: all 0s means everything belongs to batch 0
        order_mask = torch.zeros(1, device=timbre_feats.device, dtype=torch.long)
        
        # Reshape for timbre encoder: expects [N_packed, S_ref, 64]
        timbre_input = timbre_feats  # [B, S_ref, 64]
        timbre_embs, timbre_embs_mask = self.timbre_encoder(
            refer_audio_acoustic_hidden_states_packed=timbre_input,
            refer_audio_order_mask=order_mask,
            attention_mask=timbre_mask,
        )
        # timbre_embs: [B, 1, 2048] — CLS token output per batch
        
        # 4) Concatenate: [lyric, timbre, text_proj]
        # When has_timbre=0, skip timbre in the concatenation
        # For ONNX compatibility, always cat but mask the timbre contribution
        ht = has_timbre[0]
        if ht > 0:
            enc_hidden = torch.cat([lyric_out, timbre_embs, text_proj], dim=1)
        else:
            enc_hidden = torch.cat([lyric_out, text_proj], dim=1)
        
        return enc_hidden


class CondEncoderWrapperFixed(nn.Module):
    """Fixed version that always includes timbre (simplifies ONNX graph).
    
    For inference, timbre is always present (silence latent as zero timbre).
    This avoids dynamic control flow in the ONNX graph.
    
    IMPORTANT: The timbre encoder's forward() uses unpack_timbre_embeddings()
    which has data-dependent control flow (refer_audio_order_mask.max().item()).
    torch.export cannot handle this. So we manually invoke the timbre encoder's
    sub-components: embed_tokens → CLS prepend → transformer layers → norm → 
    take position 0. This is equivalent for B=1 inference.
    
    ONNX inputs:
        text_hidden:  [B, S_text, 1024]  fp16
        lyric_embed:  [B, S_lyric, 1024] fp16
        timbre_feats: [B, S_ref, 64]     fp16 (zeros if no reference)
    
    ONNX output:
        enc_hidden:   [B, S_total, 2048] fp16 where S_total = S_lyric + 1 + S_text
    """
    
    def __init__(self, cond_encoder):
        super().__init__()
        self.text_projector = cond_encoder.text_projector
        self.lyric_encoder = cond_encoder.lyric_encoder
        # Extract timbre encoder sub-components for manual invocation
        self.timbre_embed_tokens = cond_encoder.timbre_encoder.embed_tokens
        self.timbre_special_token = cond_encoder.timbre_encoder.special_token
        self.timbre_norm = cond_encoder.timbre_encoder.norm
        self.timbre_rotary_emb = cond_encoder.timbre_encoder.rotary_emb
        self.timbre_layers = cond_encoder.timbre_encoder.layers
        self.timbre_config = cond_encoder.timbre_encoder.config
    
    def _timbre_forward_simple(self, timbre_feats):
        """Run the timbre encoder without unpack_timbre_embeddings.
        
        timbre_feats: [B, S_ref, 64]
        Returns: [B, 1, hidden_size] — CLS token output
        """
        B = timbre_feats.shape[0]
        
        # Project: [B, S_ref, 64] → [B, S_ref, hidden_size]
        inputs_embeds = self.timbre_embed_tokens(timbre_feats)
        
        # Prepend CLS token: [B, S_ref+1, hidden_size]
        cls_token = self.timbre_special_token.expand(B, 1, -1)
        inputs_embeds = torch.cat([cls_token, inputs_embeds], dim=1)
        
        S = inputs_embeds.shape[1]
        
        # Position IDs and RoPE
        cache_position = torch.arange(0, S, device=inputs_embeds.device)
        position_ids = cache_position.unsqueeze(0)
        position_embeddings = self.timbre_rotary_emb(inputs_embeds, position_ids)
        
        # Build attention mask (full bidirectional, no padding)
        # Using None for SDPA = no mask = full attention
        hidden_states = inputs_embeds
        
        for layer in self.timbre_layers:
            layer_outputs = layer(
                hidden_states,
                position_embeddings,
                None,  # attention_mask=None → full bidirectional
                position_ids,
            )
            hidden_states = layer_outputs[0]
        
        hidden_states = self.timbre_norm(hidden_states)
        
        # Extract CLS token (position 0): [B, hidden_size]
        timbre_emb = hidden_states[:, 0:1, :]  # [B, 1, hidden_size]
        
        return timbre_emb
    
    def forward(self, text_hidden, lyric_embed, timbre_feats):
        B = text_hidden.shape[0]
        
        # 1) Text projection
        text_proj = self.text_projector(text_hidden)
        
        # 2) Lyric encoding
        S_lyric = lyric_embed.shape[1]
        lyric_mask = torch.ones(B, S_lyric, device=lyric_embed.device, dtype=torch.long)
        lyric_out = self.lyric_encoder(
            inputs_embeds=lyric_embed,
            attention_mask=lyric_mask,
        )
        if hasattr(lyric_out, 'last_hidden_state'):
            lyric_out = lyric_out.last_hidden_state
        else:
            lyric_out = lyric_out[0]
        
        # 3) Timbre encoding — manual path (bypasses unpack_timbre_embeddings)
        timbre_embs = self._timbre_forward_simple(timbre_feats)
        # timbre_embs: [B, 1, 2048]
        
        # 4) Concatenate: lyric + timbre + text_proj
        enc_hidden = torch.cat([lyric_out, timbre_embs, text_proj], dim=1)
        
        return enc_hidden


def load_model(model_dir: str, device: str = "cuda", dtype=torch.float32):
    """Load the AceStep model and extract the condition encoder."""
    model_dir = Path(model_dir)
    
    if sys.platform == "win32":
        import io
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    
    # Monkey-patch transformers auto_docstring
    try:
        import transformers.utils.auto_docstring as _ad
        _ad.auto_docstring = lambda *a, **kw: (lambda cls: cls)
    except Exception:
        pass
    
    sys.path.insert(0, str(model_dir))
    
    # Create a stub AceStepConfig module to bypass the 'acestep' package import.
    # The real AceStepConfig is a PretrainedConfig subclass. We construct it
    # using AutoConfig which reads config.json and finds the auto_map.
    # But first we need the config module to exist so the model code can import it.
    import json
    import types
    from transformers import PretrainedConfig
    
    with open(model_dir / "config.json") as f:
        config_dict = json.load(f)
    
    # Create AceStepConfig class dynamically from config.json
    class AceStepConfig(PretrainedConfig):
        model_type = "acestep"
        
        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            # Set all config keys as attributes
            for k, v in kwargs.items():
                if not hasattr(self, k):
                    setattr(self, k, v)
            # Ensure critical attrs have defaults
            if not hasattr(self, 'text_hidden_dim'):
                self.text_hidden_dim = 1024
            if not hasattr(self, 'timbre_hidden_dim'):
                self.timbre_hidden_dim = 64
            if not hasattr(self, 'encoder_hidden_size'):
                self.encoder_hidden_size = 2048
            if not hasattr(self, 'encoder_intermediate_size'):
                self.encoder_intermediate_size = 6144
            if not hasattr(self, 'encoder_num_attention_heads'):
                self.encoder_num_attention_heads = 16
            if not hasattr(self, 'encoder_num_key_value_heads'):
                self.encoder_num_key_value_heads = 8
            if not hasattr(self, 'num_lyric_encoder_hidden_layers'):
                self.num_lyric_encoder_hidden_layers = 8
            if not hasattr(self, 'num_timbre_encoder_hidden_layers'):
                self.num_timbre_encoder_hidden_layers = 4
            if not hasattr(self, 'num_attention_pooler_hidden_layers'):
                self.num_attention_pooler_hidden_layers = 2
            if not hasattr(self, 'out_channels'):
                self.out_channels = 64
            if not hasattr(self, 'in_channels'):
                self.in_channels = 192
    
    # Register the stub config module
    stub_mod = types.ModuleType("configuration_acestep_v15")
    stub_mod.AceStepConfig = AceStepConfig
    sys.modules["configuration_acestep_v15"] = stub_mod
    
    # Create stub acestep package hierarchy. Each intermediate module needs
    # __path__ set so it acts as a package (allows submodule imports).
    for mod_name in ["acestep", "acestep.models", "acestep.models.common"]:
        m = types.ModuleType(mod_name)
        m.__path__ = []  # makes it act as a package
        sys.modules[mod_name] = m
    
    # Register configuration_acestep_v15 under the acestep.models.common path
    cfg_mod = types.ModuleType("acestep.models.common.configuration_acestep_v15")
    cfg_mod.AceStepConfig = AceStepConfig
    sys.modules["acestep.models.common.configuration_acestep_v15"] = cfg_mod
    
    # Create apg_guidance stub — these functions are used by the DiT diffusion
    # loop but NOT by the condition encoder. Provide dummies to satisfy import.
    class MomentumBuffer:
        def __init__(self, *a, **kw): pass
    def _apg_stub(*a, **kw): return None
    
    apg_stub = types.ModuleType("acestep.models.common.apg_guidance")
    apg_stub.MomentumBuffer = MomentumBuffer
    apg_stub.adg_forward = _apg_stub
    apg_stub.adg_w_norm_forward = _apg_stub
    apg_stub.adg_wo_clip_forward = _apg_stub
    apg_stub.apg_forward = _apg_stub
    apg_stub.cfg_forward = _apg_stub
    apg_stub.call_cos_tensor = _apg_stub
    apg_stub.compute_perpendicular_component = _apg_stub
    apg_stub.project = _apg_stub
    sys.modules["acestep.models.common.apg_guidance"] = apg_stub
    
    # Also register the local apg_guidance module
    apg_local = types.ModuleType("apg_guidance")
    apg_local.MomentumBuffer = MomentumBuffer
    apg_local.adg_forward = _apg_stub
    apg_local.apg_forward = _apg_stub
    apg_local.cfg_forward = _apg_stub
    sys.modules["apg_guidance"] = apg_local
    
    config = AceStepConfig(**config_dict)
    config._attn_implementation = "sdpa"
    
    # The encoder's Qwen3 sub-models (lyric/timbre) use encoder_hidden_size
    # as their hidden_size. Set it on the config so Qwen3RotaryEmbedding works.
    # transformers 5.x requires rope_parameters dict.
    if not hasattr(config, 'rope_parameters') or config.rope_parameters is None:
        config.rope_parameters = {
            "rope_type": "default",
            "rope_theta": config.rope_theta if hasattr(config, 'rope_theta') else 1000000.0,
        }
    
    print(f"[export_cond_enc] Loading model from {model_dir}...")
    t0 = time.time()
    
    # The full model creates a separate encoder config with encoder-specific
    # dimensions (see AceStepConditionGenerationModel.__init__ lines 1621-1628).
    # The encoder uses encoder_hidden_size (2048), not the DiT hidden_size (2560).
    import copy
    encoder_config = copy.deepcopy(config)
    encoder_config.hidden_size = config.encoder_hidden_size
    encoder_config.intermediate_size = config.encoder_intermediate_size
    encoder_config.num_attention_heads = config.encoder_num_attention_heads
    encoder_config.num_key_value_heads = config.encoder_num_key_value_heads
    
    from modeling_acestep_v15_xl_base import AceStepConditionEncoder
    
    cond_encoder = AceStepConditionEncoder(encoder_config)
    
    # Load weights — filter to encoder.* prefix
    from safetensors.torch import load_file
    st_path = model_dir / "model.safetensors"
    state_dict = load_file(str(st_path))
    
    cond_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith("encoder."):
            cond_state_dict[k[len("encoder."):]] = v
    
    missing, unexpected = cond_encoder.load_state_dict(cond_state_dict, strict=False)
    if missing:
        print(f"[export_cond_enc] Warning: {len(missing)} missing keys (first 5: {missing[:5]})")
    if unexpected:
        print(f"[export_cond_enc] Warning: {len(unexpected)} unexpected keys (first 5: {unexpected[:5]})")
    
    cond_encoder = cond_encoder.to(device=device, dtype=dtype)
    cond_encoder.eval()
    
    t1 = time.time()
    n_params = sum(p.numel() for p in cond_encoder.parameters()) / 1e6
    print(f"[export_cond_enc] Model loaded in {t1-t0:.1f}s ({n_params:.0f}M params)")
    print(f"[export_cond_enc] text_hidden_dim={encoder_config.text_hidden_dim}, hidden_size={encoder_config.hidden_size}")
    
    return cond_encoder, encoder_config


def export_onnx(cond_encoder, config, output_path: str, opset: int = 18):
    """Export the condition encoder to ONNX."""
    device = next(cond_encoder.parameters()).device
    dtype = next(cond_encoder.parameters()).dtype
    
    wrapper = CondEncoderWrapperFixed(cond_encoder)
    wrapper.eval()
    
    # Dummy inputs
    B = 1
    S_text = 64
    S_lyric = 128
    S_ref = 8   # 8 frames of reference audio (short clip)
    
    dummy_text_hidden = torch.randn(B, S_text, config.text_hidden_dim, device=device, dtype=dtype)
    dummy_lyric_embed = torch.randn(B, S_lyric, config.text_hidden_dim, device=device, dtype=dtype)
    dummy_timbre_feats = torch.randn(B, S_ref, config.timbre_hidden_dim, device=device, dtype=dtype)
    
    print(f"[export_cond_enc] Tracing with shapes: text={list(dummy_text_hidden.shape)}, "
          f"lyric={list(dummy_lyric_embed.shape)}, timbre={list(dummy_timbre_feats.shape)}")
    
    # Test forward
    print("[export_cond_enc] Testing forward pass...")
    with torch.no_grad():
        test_out = wrapper(dummy_text_hidden, dummy_lyric_embed, dummy_timbre_feats)
    expected_S = S_lyric + 1 + S_text  # lyric + timbre(1) + text
    print(f"[export_cond_enc] Output shape: {list(test_out.shape)} "
          f"(expected [{B}, {expected_S}, {config.hidden_size}])")
    
    # Export
    print(f"[export_cond_enc] Exporting to ONNX (opset {opset})...")
    t0 = time.time()
    
    torch.onnx.export(
        wrapper,
        (dummy_text_hidden, dummy_lyric_embed, dummy_timbre_feats),
        output_path,
        opset_version=opset,
        input_names=["text_hidden", "lyric_embed", "timbre_feats"],
        output_names=["enc_hidden"],
        dynamic_axes={
            "text_hidden":  {0: "batch", 1: "text_seq"},
            "lyric_embed":  {0: "batch", 1: "lyric_seq"},
            "timbre_feats": {0: "batch", 1: "timbre_seq"},
            "enc_hidden":   {0: "batch", 1: "enc_seq"},
        },
        do_constant_folding=True,
        export_params=True,
    )
    
    t1 = time.time()
    file_size = os.path.getsize(output_path)
    print(f"[export_cond_enc] Exported to {output_path}")
    print(f"[export_cond_enc] File size: {file_size/1e6:.1f} MB")
    print(f"[export_cond_enc] Export time: {t1-t0:.1f}s")


def export_null_cond_emb(model_dir: str, output_path: str):
    """Export null_condition_emb as raw float32 binary."""
    from safetensors.torch import load_file
    model_dir = Path(model_dir)
    st_path = model_dir / "model.safetensors"
    
    state_dict = load_file(str(st_path))
    key = "null_condition_emb"
    if key not in state_dict:
        print(f"[export_cond_enc] WARNING: {key} not found, skipping")
        return
    
    vec = state_dict[key].detach().cpu().float().numpy().flatten()
    with open(output_path, "wb") as f:
        f.write(struct.pack("<I", len(vec)))
        f.write(vec.tobytes())
    
    print(f"[export_cond_enc] null_condition_emb: [{len(vec)}] -> {output_path} ({len(vec)*4} bytes)")


def verify_onnx(onnx_path: str, cond_encoder, config):
    """Verify ONNX output matches PyTorch."""
    try:
        import onnxruntime as ort
    except ImportError:
        print("[export_cond_enc] onnxruntime not installed, skipping verification")
        return
    
    device = next(cond_encoder.parameters()).device
    dtype = next(cond_encoder.parameters()).dtype
    wrapper = CondEncoderWrapperFixed(cond_encoder)
    wrapper.eval()
    
    B, S_text, S_lyric, S_ref = 1, 32, 64, 8
    text_hidden = torch.randn(B, S_text, config.text_hidden_dim, device=device, dtype=dtype)
    lyric_embed = torch.randn(B, S_lyric, config.text_hidden_dim, device=device, dtype=dtype)
    timbre_feats = torch.randn(B, S_ref, config.timbre_hidden_dim, device=device, dtype=dtype)
    
    with torch.no_grad():
        ref_out = wrapper(text_hidden, lyric_embed, timbre_feats).cpu().float().numpy()
    
    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    sess = ort.InferenceSession(onnx_path, providers=providers)
    ort_out = sess.run(None, {
        "text_hidden": text_hidden.cpu().float().numpy(),
        "lyric_embed": lyric_embed.cpu().float().numpy(),
        "timbre_feats": timbre_feats.cpu().float().numpy(),
    })[0]
    
    max_diff = np.max(np.abs(ref_out - ort_out))
    mean_diff = np.mean(np.abs(ref_out - ort_out))
    print(f"[export_cond_enc] Verification: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}")
    
    if max_diff < 0.05:
        print("[export_cond_enc] PASS: ONNX output matches PyTorch")
    else:
        print("[export_cond_enc] WARNING: Large difference — may need investigation")


def main():
    parser = argparse.ArgumentParser(description="Export AceStep condition encoder to ONNX")
    parser.add_argument("--model-dir", required=True,
                        help="Path to the DiT model directory (contains encoder weights)")
    parser.add_argument("--output", default=None,
                        help="Output ONNX file (default: models/onnx/cond_encoder.onnx)")
    parser.add_argument("--opset", type=int, default=18)
    parser.add_argument("--verify", action="store_true")
    parser.add_argument("--device", default="cuda")
    args = parser.parse_args()
    
    if args.output is None:
        onnx_dir = Path(args.model_dir).parent / "onnx"
        onnx_dir.mkdir(parents=True, exist_ok=True)
        args.output = str(onnx_dir / "cond_encoder.onnx")
    
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    output_dir = os.path.dirname(args.output)
    
    # Load model
    cond_encoder, config = load_model(args.model_dir, device=args.device)
    
    # Export ONNX
    export_onnx(cond_encoder, config, args.output, opset=args.opset)
    
    # Export null_condition_emb
    null_cond_path = os.path.join(output_dir, "null_condition_emb.bin")
    export_null_cond_emb(args.model_dir, null_cond_path)
    
    # Verify
    if args.verify:
        verify_onnx(args.output, cond_encoder, config)
    
    print("[export_cond_enc] Done!")


if __name__ == "__main__":
    main()