#!/usr/bin/env python3
"""Export FastOobleckDecoder checkpoint to ONNX and build a TRT engine.

The student decoder has the exact same I/O contract as the teacher
(latents [B, 64, T] -> audio [B, 2, samples]), so we reuse the existing
VAE TRT build pipeline. The resulting engine is a drop-in replacement.

Usage:
    uv run python research_program/vae_distillation/export_trt.py \
        --ckpt research_program/vae_distillation/checkpoints/student_step620000.pt

    # FP32 engine (if FP16 has Snake1d issues):
    uv run python research_program/vae_distillation/export_trt.py \
        --ckpt research_program/vae_distillation/checkpoints/student_step620000.pt \
        --no-fp16

    # Custom output directory:
    uv run python research_program/vae_distillation/export_trt.py \
        --ckpt ... --output-dir trt_engines
"""

import argparse
import logging
import math
import sys
from pathlib import Path

import torch
import torch.nn as nn
from torch.nn.utils import weight_norm

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


# =========================================================================
# Student model definition (must match training script exactly)
# =========================================================================

class Snake1d(nn.Module):
    def __init__(self, hidden_dim, logscale=True):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1))
        self.alpha.requires_grad = True
        self.beta.requires_grad = True
        self.logscale = logscale

    def forward(self, hidden_states):
        shape = hidden_states.shape
        alpha = self.alpha if not self.logscale else torch.exp(self.alpha)
        beta = self.beta if not self.logscale else torch.exp(self.beta)
        hidden_states = hidden_states.reshape(shape[0], shape[1], -1)
        hidden_states = hidden_states + (beta + 1e-9).reciprocal() * torch.sin(alpha * hidden_states).pow(2)
        hidden_states = hidden_states.reshape(shape)
        return hidden_states


class FastResidualUnit(nn.Module):
    def __init__(self, dim: int, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.snake1 = Snake1d(dim)
        self.conv1 = weight_norm(nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad))
        self.snake2 = Snake1d(dim)
        self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1))

    def forward(self, x):
        h = self.conv1(self.snake1(x))
        h = self.conv2(self.snake2(h))
        pad = (x.shape[-1] - h.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + h


class FastDecoderBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, stride: int = 1):
        super().__init__()
        self.snake1 = Snake1d(in_dim)
        self.conv_t = weight_norm(nn.ConvTranspose1d(
            in_dim, out_dim, kernel_size=2 * stride, stride=stride,
            padding=math.ceil(stride / 2),
        ))
        self.res1 = FastResidualUnit(out_dim, dilation=1)
        self.res2 = FastResidualUnit(out_dim, dilation=3)

    def forward(self, x):
        x = self.snake1(x)
        x = self.conv_t(x)
        x = self.res1(x)
        x = self.res2(x)
        return x


class FastOobleckDecoder(nn.Module):
    def __init__(self, channels=128, input_channels=64, audio_channels=2,
                 upsampling_ratios=None, channel_multiples=None):
        super().__init__()
        if upsampling_ratios is None:
            upsampling_ratios = [10, 6, 4, 4, 2]
        if channel_multiples is None:
            channel_multiples = [1, 2, 4, 8, 8]
        cm = [1] + channel_multiples
        self.conv1 = weight_norm(nn.Conv1d(input_channels, channels * cm[-1], kernel_size=7, padding=3))
        blocks = []
        for i, stride in enumerate(upsampling_ratios):
            in_dim = channels * cm[len(upsampling_ratios) - i]
            out_dim = channels * cm[len(upsampling_ratios) - i - 1]
            blocks.append(FastDecoderBlock(in_dim, out_dim, stride=stride))
        self.blocks = nn.ModuleList(blocks)
        self.final_snake = Snake1d(channels)
        self.conv2 = weight_norm(nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False))

    def forward(self, latents: torch.Tensor) -> torch.Tensor:
        x = self.conv1(latents)
        for block in self.blocks:
            x = block(x)
        x = self.final_snake(x)
        x = self.conv2(x)
        return x


# =========================================================================
# Export
# =========================================================================

def export_onnx(student: nn.Module, onnx_path: Path, device: str = "cuda"):
    """Export student decoder to ONNX with dynamic latent length."""
    onnx_path.parent.mkdir(parents=True, exist_ok=True)

    # Trace with 30s worth of latent frames (750 = 30s * 25 fps)
    example = torch.randn(1, 64, 750, device=device, dtype=torch.float32)

    logger.info("Tracing student decoder for ONNX export...")
    with torch.no_grad():
        torch.onnx.export(
            student,
            (example,),
            str(onnx_path),
            input_names=["latents"],
            output_names=["audio"],
            dynamic_axes={
                "latents": {0: "batch", 2: "latent_frames"},
                "audio": {0: "batch", 2: "samples"},
            },
            opset_version=18,
            do_constant_folding=True,
            dynamo=False,
        )
    logger.info("ONNX saved to %s (%.1f MB)", onnx_path, onnx_path.stat().st_size / (1 << 20))
    return onnx_path


def build_engine(onnx_path: Path, engine_path: Path, fp16: bool = True):
    """Build TRT engine using the same config as the teacher VAE decoder."""
    # Import the existing build infrastructure
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from acestep.engine.trt.vae_export import VAETRTBuildConfig, build_vae_decode_engine

    config = VAETRTBuildConfig(fp16=fp16)
    return build_vae_decode_engine(str(onnx_path), str(engine_path), config=config)


def verify_engine(engine_path: Path, device: str = "cuda"):
    """Quick sanity check: run a dummy input through the engine."""
    from polygraphy.backend.common import bytes_from_path
    from polygraphy.backend.trt import engine_from_bytes
    from polygraphy import cuda as pg_cuda

    engine = engine_from_bytes(bytes_from_path(str(engine_path)))
    ctx = engine.create_execution_context()
    stream = pg_cuda.Stream()

    # 10s of audio = 250 latent frames
    T = 250
    latents = torch.randn(1, 64, T, device=device, dtype=torch.float32).contiguous()
    ctx.set_input_shape("latents", (1, 64, T))
    ctx.set_tensor_address("latents", latents.data_ptr())

    out_shape = tuple(ctx.get_tensor_shape("audio"))
    audio_buf = torch.empty(out_shape, dtype=torch.float32, device=device)
    ctx.set_tensor_address("audio", audio_buf.data_ptr())

    ok = ctx.execute_async_v3(stream.ptr)
    stream.synchronize()

    expected_samples = T * 1920  # hop length
    assert ok, "TRT execution failed"
    assert out_shape == (1, 2, expected_samples), f"Shape mismatch: {out_shape} vs (1, 2, {expected_samples})"
    assert not torch.isnan(audio_buf).any(), "NaN in output"
    assert audio_buf.abs().max() > 1e-6, "Output is all zeros"

    logger.info("Engine verification passed: input [1, 64, %d] -> output %s, range [%.4f, %.4f]",
                T, list(out_shape), audio_buf.min().item(), audio_buf.max().item())

    # Quick speed check
    import time
    torch.cuda.synchronize()
    times = []
    for _ in range(20):
        torch.cuda.synchronize()
        t0 = time.time()
        ctx.execute_async_v3(stream.ptr)
        stream.synchronize()
        times.append(time.time() - t0)
    avg_ms = sum(times) / len(times) * 1000
    logger.info("TRT speed (10s audio, 20 trials): %.1f ms avg", avg_ms)


def main():
    parser = argparse.ArgumentParser(description="Export FastOobleckDecoder to ONNX + TRT")
    parser.add_argument("--ckpt", type=str, required=True, help="Path to student checkpoint .pt")
    parser.add_argument("--output-dir", type=str, default=None,
                        help="Output directory (default: trt_engines/ in project root)")
    parser.add_argument("--no-fp16", action="store_true", help="Build FP32 engine instead of FP16")
    parser.add_argument("--skip-engine", action="store_true", help="Only export ONNX, skip TRT build")
    parser.add_argument("--skip-verify", action="store_true", help="Skip engine verification")
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args()

    ckpt_path = Path(args.ckpt)
    step = "unknown"

    # Determine output paths
    project_root = Path(__file__).resolve().parents[2]
    if args.output_dir:
        out_dir = Path(args.output_dir)
    else:
        out_dir = project_root / "trt_engines"

    # Load student
    logger.info("Loading checkpoint: %s", ckpt_path)
    ckpt = torch.load(ckpt_path, map_location=args.device, weights_only=False)

    config = ckpt.get("config", {})
    student = FastOobleckDecoder(
        channels=config.get("channels", 128),
        input_channels=config.get("input_channels", 64),
        audio_channels=config.get("audio_channels", 2),
        upsampling_ratios=config.get("upsampling_ratios", [10, 6, 4, 4, 2]),
        channel_multiples=config.get("channel_multiples", [1, 2, 4, 8, 8]),
    ).to(args.device).eval()
    student.load_state_dict(ckpt["student_state_dict"])
    step = ckpt.get("step", "unknown")
    params_m = sum(p.numel() for p in student.parameters()) / 1e6
    logger.info("Student loaded: step %s, %.1fM params", step, params_m)

    # ONNX export
    onnx_dir = out_dir / "_onnx" / "dreamvae_decode"
    onnx_path = onnx_dir / "dreamvae_decode.onnx"
    export_onnx(student, onnx_path, device=args.device)

    if args.skip_engine:
        logger.info("Skipping TRT engine build (--skip-engine)")
        return

    # TRT engine build
    fp16 = not args.no_fp16
    prec = "fp16" if fp16 else "fp32"
    engine_name = f"dreamvae_decode_{prec}_240s"
    engine_dir = out_dir / engine_name
    engine_path = engine_dir / f"{engine_name}.engine"

    logger.info("Building TRT engine (%s)...", prec)
    build_engine(onnx_path, engine_path, fp16=fp16)

    if not args.skip_verify:
        logger.info("Verifying engine...")
        verify_engine(engine_path, device=args.device)

    logger.info("Done. Engine: %s", engine_path)
    logger.info("To use as drop-in replacement, pass this engine path where vae_decode engine is expected.")


if __name__ == "__main__":
    main()