#!/usr/bin/env python3
"""Export Meta's HTDemucs model to ONNX format for use with RustyTorch++.

Usage:
    pip install torch demucs onnx
    python export_demucs_onnx.py --stems 4 --output models/htdemucs.onnx
    python export_demucs_onnx.py --stems 6 --output models/htdemucs_6s.onnx

The exported model expects:
    Input:  "mix"    shape [1, 2, segment_length]  (stereo audio at 44100 Hz)
    Output: "stems"  shape [1, num_stems, 2, segment_length]

Segment length is fixed at export time (default: 441000 = 10 seconds at 44.1 kHz).
The Rust inference code handles longer audio via segmented overlap-add.
"""

import argparse
import os

import torch
import torch.onnx


def export_demucs(num_stems: int, output_path: str, segment_length: int = 441000):
    # Import demucs and load the pretrained model
    from demucs.pretrained import get_model

    model_name = "htdemucs" if num_stems == 4 else "htdemucs_6s"
    print(f"Loading {model_name} ({num_stems} stems)...")

    model = get_model(model_name)
    model.eval()

    # Create dummy input: [batch=1, channels=2, samples=segment_length]
    dummy_input = torch.randn(1, 2, segment_length)

    print(f"Exporting to ONNX: {output_path}")
    print(f"  Input shape:  [1, 2, {segment_length}]")
    print(f"  Output shape: [1, {num_stems}, 2, {segment_length}]")

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

    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        input_names=["mix"],
        output_names=["stems"],
        opset_version=17,
        do_constant_folding=True,
        dynamic_axes=None,  # Fixed shapes for reliable ONNX inference
    )

    # Verify the exported model
    import onnx
    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)

    file_size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"Export complete: {output_path} ({file_size_mb:.1f} MB)")
    print(f"Model verified successfully.")


def main():
    parser = argparse.ArgumentParser(description="Export HTDemucs to ONNX")
    parser.add_argument("--stems", type=int, default=4, choices=[4, 6],
                        help="Number of stems (4 or 6)")
    parser.add_argument("--output", type=str, default=None,
                        help="Output ONNX file path")
    parser.add_argument("--segment-length", type=int, default=441000,
                        help="Segment length in samples (default: 441000 = 10s at 44.1kHz)")
    args = parser.parse_args()

    if args.output is None:
        name = "htdemucs" if args.stems == 4 else "htdemucs_6s"
        args.output = f"models/{name}.onnx"

    export_demucs(args.stems, args.output, args.segment_length)


if __name__ == "__main__":
    main()