#!/usr/bin/env python3
"""Export PP-VAE (LeVo autoencoder_music_1320k.ckpt) to ONNX format.

Exports BOTH encoder and decoder as separate ONNX files for use with
TensorRT or ONNX Runtime in the HOT-Step-CPP engine.

PP-VAE is a different model from scragvae (same Oobleck architecture,
different weights). Source: tencent-ailab/SongGeneration autoencoder_music_1320k

Tensor specs:
  Encoder:
    Input:  "audio"    [B, 2, T_audio]   (stereo 48kHz)
    Output: "latents"  [B, 64, T_latent] (mean-only, deterministic)

  Decoder:
    Input:  "latents"  [B, 64, T_latent] (latent channels @ 25Hz)
    Output: "audio"    [B, 2, T_audio]   (stereo, T_audio = T_latent * 1920)

Usage:
  python export_pp_vae.py --ckpt path/to/autoencoder_music_1320k.ckpt --output-dir models/onnx/
  python export_pp_vae.py  # uses default paths
"""

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

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


# ─── Model Architecture (from LeVo / stable-audio) ──────────────────────────
# Matches the architecture in levo-vae/reencode.py exactly.

class Snake1d(nn.Module):
    """Snake activation: y = x + sin²(exp(α)·x) / exp(β)"""
    def __init__(self, channels):
        super().__init__()
        self.alpha = nn.Parameter(torch.ones(channels))
        self.beta = nn.Parameter(torch.ones(channels))

    def forward(self, x):
        a = torch.exp(self.alpha).unsqueeze(0).unsqueeze(-1)
        b = torch.exp(self.beta).unsqueeze(0).unsqueeze(-1)
        return x + (torch.sin(a * x) ** 2) / (b + 1e-9)


def WNConv1d(*args, **kwargs):
    return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs))

def WNConvTranspose1d(*args, **kwargs):
    return nn.utils.weight_norm(nn.ConvTranspose1d(*args, **kwargs))


class ResUnit(nn.Module):
    """Residual unit: snake → dilated conv(k=7) → snake → conv(k=1) → + skip"""
    def __init__(self, channels, dilation):
        super().__init__()
        self.layers = nn.Sequential(
            Snake1d(channels),
            WNConv1d(channels, channels, kernel_size=7, dilation=dilation,
                     padding=3 * dilation),
            Snake1d(channels),
            WNConv1d(channels, channels, kernel_size=1),
        )

    def forward(self, x):
        return x + self.layers(x)


class EncoderBlock(nn.Module):
    """3× ResUnit → Snake → strided Conv1d (downsample)"""
    def __init__(self, in_ch, out_ch, stride):
        super().__init__()
        layers = []
        for dil in [1, 3, 9]:
            layers.append(ResUnit(in_ch, dil))
        layers.append(Snake1d(in_ch))
        layers.append(WNConv1d(in_ch, out_ch, kernel_size=stride * 2,
                               stride=stride, padding=stride // 2))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class DecoderBlock(nn.Module):
    """Snake → ConvTranspose1d (upsample) → 3× ResUnit"""
    def __init__(self, in_ch, out_ch, stride):
        super().__init__()
        layers = []
        layers.append(Snake1d(in_ch))
        layers.append(WNConvTranspose1d(in_ch, out_ch, kernel_size=stride * 2,
                                         stride=stride, padding=stride // 2))
        for dil in [1, 3, 9]:
            layers.append(ResUnit(out_ch, dil))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class OobleckEncoder(nn.Module):
    def __init__(self, in_channels=2, channels=128, c_mults=[1,2,4,8,16],
                 strides=[2,4,4,6,10], latent_dim=128, **kw):
        super().__init__()
        c_mults = [1] + c_mults
        layers = [WNConv1d(in_channels, channels * c_mults[0], kernel_size=7, padding=3)]
        for i, stride in enumerate(strides):
            layers.append(EncoderBlock(channels * c_mults[i], channels * c_mults[i + 1], stride))
        layers.append(Snake1d(channels * c_mults[-1]))
        layers.append(WNConv1d(channels * c_mults[-1], latent_dim, kernel_size=3, padding=1))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


class OobleckDecoder(nn.Module):
    def __init__(self, out_channels=2, channels=128, c_mults=[1,2,4,8,16],
                 strides=[2,4,4,6,10], latent_dim=64, **kw):
        super().__init__()
        c_mults = [1] + c_mults
        c_mults_rev = list(reversed(c_mults))
        strides_rev = list(reversed(strides))
        layers = [WNConv1d(latent_dim, channels * c_mults_rev[0], kernel_size=7, padding=3)]
        for i, stride in enumerate(strides_rev):
            layers.append(DecoderBlock(channels * c_mults_rev[i], channels * c_mults_rev[i + 1], stride))
        layers.append(Snake1d(channels * c_mults_rev[-1]))
        layers.append(WNConv1d(channels * c_mults_rev[-1], out_channels, kernel_size=7, padding=3, bias=False))
        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        return self.layers(x)


# ─── Encoder Wrapper (deterministic, mean-only) ─────────────────────────────

class PPVAEEncoderWrapper(nn.Module):
    """Wraps OobleckEncoder to output only the mean (first 64 of 128 channels).

    The encoder outputs 128 channels: [mean(64), logvar(64)].
    For deterministic encoding we only need the mean.
    """
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder

    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        h = self.encoder(audio)  # [B, 128, T_latent]
        mean = h[:, :64, :]     # [B, 64, T_latent]
        return mean


# ─── Export Functions ────────────────────────────────────────────────────────

def export_encoder(ckpt_path: str, output_path: str, opset: int = 18) -> str:
    """Export PP-VAE encoder to ONNX (deterministic, mean-only)."""
    print(f"\n{'='*60}")
    print(f"  Exporting PP-VAE Encoder")
    print(f"{'='*60}")

    print(f"Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt.get("state_dict", ckpt)
    del ckpt

    enc_cfg = {"in_channels": 2, "channels": 128, "c_mults": [1,2,4,8,16],
               "strides": [2,4,4,6,10], "latent_dim": 128}

    encoder = OobleckEncoder(**enc_cfg)
    enc_sd = {k.replace("encoder.", ""): v for k, v in sd.items() if k.startswith("encoder.")}
    encoder.load_state_dict(enc_sd)
    del enc_sd, sd

    wrapper = PPVAEEncoderWrapper(encoder)
    wrapper.eval()

    # Dummy: [1, 2, 10s * 48kHz]
    T_audio = 250 * 1920  # 250 latent frames = 10s
    dummy = torch.randn(1, 2, T_audio, dtype=torch.float32)
    print(f"Dummy input: {dummy.shape}")

    with torch.no_grad():
        test_out = wrapper(dummy)
    print(f"Test output: {test_out.shape} (expected [1, 64, 250])")

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    print(f"Exporting to ONNX (opset {opset})...")
    t0 = time.time()

    torch.onnx.export(
        wrapper,
        (dummy,),
        output_path,
        opset_version=opset,
        input_names=["audio"],
        output_names=["latents"],
        dynamic_axes={
            "audio": {0: "batch", 2: "samples"},
            "latents": {0: "batch", 2: "latent_frames"},
        },
        do_constant_folding=True,
    )

    elapsed = time.time() - t0
    size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"Exported in {elapsed:.1f}s: {output_path} ({size_mb:.1f} MB)")

    # Validate
    _validate_encoder(output_path, wrapper, dummy)

    del wrapper, encoder
    gc.collect()
    return output_path


def export_decoder(ckpt_path: str, output_path: str, opset: int = 18) -> str:
    """Export PP-VAE decoder to ONNX."""
    print(f"\n{'='*60}")
    print(f"  Exporting PP-VAE Decoder")
    print(f"{'='*60}")

    print(f"Loading checkpoint: {ckpt_path}")
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt.get("state_dict", ckpt)
    del ckpt

    dec_cfg = {"out_channels": 2, "channels": 128, "c_mults": [1,2,4,8,16],
               "strides": [2,4,4,6,10], "latent_dim": 64}

    decoder = OobleckDecoder(**dec_cfg)
    dec_sd = {k.replace("decoder.", ""): v for k, v in sd.items() if k.startswith("decoder.")}
    decoder.load_state_dict(dec_sd)
    del dec_sd, sd

    decoder.eval()

    # Dummy: [1, 64, 250] = 10s of latents
    dummy = torch.randn(1, 64, 250, dtype=torch.float32)
    print(f"Dummy input: {dummy.shape}")

    with torch.no_grad():
        test_out = decoder(dummy)
    print(f"Test output: {test_out.shape} (expected [1, 2, {250 * 1920}])")

    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    print(f"Exporting to ONNX (opset {opset})...")
    t0 = time.time()

    torch.onnx.export(
        decoder,
        (dummy,),
        output_path,
        opset_version=opset,
        input_names=["latents"],
        output_names=["audio"],
        dynamic_axes={
            "latents": {0: "batch", 2: "latent_frames"},
            "audio": {0: "batch", 2: "samples"},
        },
        do_constant_folding=True,
    )

    elapsed = time.time() - t0
    size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"Exported in {elapsed:.1f}s: {output_path} ({size_mb:.1f} MB)")

    # Validate
    _validate_decoder(output_path, decoder, dummy)

    del decoder
    gc.collect()
    return output_path


# ─── Validation ──────────────────────────────────────────────────────────────

def _validate_encoder(onnx_path, pytorch_model, dummy_input):
    """Compare ONNX encoder output against PyTorch reference."""
    try:
        import onnxruntime as ort
    except ImportError:
        print("[WARN] onnxruntime not available, skipping validation")
        return

    print("\nValidating encoder ONNX vs PyTorch...")
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

    with torch.no_grad():
        ref = pytorch_model(dummy_input).numpy()

    ort_out = sess.run(None, {"audio": dummy_input.numpy()})[0]

    diff = np.abs(ref - ort_out)
    print(f"  Max diff: {diff.max():.6f}")
    print(f"  Mean diff: {diff.mean():.8f}")
    print(f"  Shape match: {ref.shape == ort_out.shape}")

    if diff.max() < 0.01:
        print("  ✓ PASS")
    else:
        print("  ✗ FAIL — large deviation!")


def _validate_decoder(onnx_path, pytorch_model, dummy_input):
    """Compare ONNX decoder output against PyTorch reference."""
    try:
        import onnxruntime as ort
    except ImportError:
        print("[WARN] onnxruntime not available, skipping validation")
        return

    print("\nValidating decoder ONNX vs PyTorch...")
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])

    with torch.no_grad():
        ref = pytorch_model(dummy_input).numpy()

    ort_out = sess.run(None, {"latents": dummy_input.numpy()})[0]

    diff = np.abs(ref - ort_out)
    print(f"  Max diff: {diff.max():.6f}")
    print(f"  Mean diff: {diff.mean():.8f}")
    print(f"  Shape match: {ref.shape == ort_out.shape}")

    if diff.max() < 0.01:
        print("  ✓ PASS")
    else:
        print("  ✗ FAIL — large deviation!")


# ─── CLI ─────────────────────────────────────────────────────────────────────

def main():
    default_ckpt = r"D:\Ace-Step-Latest\levo-vae\autoencoder_music_1320k.ckpt"
    default_output_dir = r"D:\Ace-Step-Latest\hot-step-cpp\models\onnx"

    parser = argparse.ArgumentParser(description="Export PP-VAE to ONNX (encoder + decoder)")
    parser.add_argument("--ckpt", default=default_ckpt,
                        help=f"Path to autoencoder_music_1320k.ckpt (default: {default_ckpt})")
    parser.add_argument("--output-dir", default=default_output_dir,
                        help=f"Output directory for ONNX files (default: {default_output_dir})")
    parser.add_argument("--opset", type=int, default=18,
                        help="ONNX opset version (default: 18)")
    parser.add_argument("--encoder-only", action="store_true",
                        help="Export only the encoder")
    parser.add_argument("--decoder-only", action="store_true",
                        help="Export only the decoder")
    args = parser.parse_args()

    if not os.path.exists(args.ckpt):
        print(f"[ERROR] Checkpoint not found: {args.ckpt}")
        sys.exit(1)

    os.makedirs(args.output_dir, exist_ok=True)

    enc_path = os.path.join(args.output_dir, "pp-vae_encoder.onnx")
    dec_path = os.path.join(args.output_dir, "pp-vae_decoder.onnx")

    if not args.decoder_only:
        export_encoder(args.ckpt, enc_path, args.opset)

    if not args.encoder_only:
        export_decoder(args.ckpt, dec_path, args.opset)

    print(f"\n{'='*60}")
    print(f"  Done!")
    if not args.decoder_only:
        print(f"  Encoder: {enc_path}")
    if not args.encoder_only:
        print(f"  Decoder: {dec_path}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()