#!/usr/bin/env python3
"""Export MetricGAN enhance_model (BLSTM masker) to ONNX. Uses project .venv."""
from __future__ import annotations

import argparse
from pathlib import Path

import torch
import torch.nn as nn
from speechbrain.inference.enhancement import SpectralMaskEnhancement

SAVEDIR = "/root/.cache/speechbrain-enhancement"


class MaskNet(nn.Module):
    def __init__(self, net: nn.Module):
        super().__init__()
        self.net = net

    def forward(self, feats: torch.Tensor) -> torch.Tensor:
        # lengths=1.0 relative, one per batch row
        lengths = torch.ones(feats.shape[0], dtype=feats.dtype)
        return self.net(feats, lengths=lengths)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="scripts/metricgan_enhance.onnx")
    args = ap.parse_args()
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)

    m = SpectralMaskEnhancement.from_hparams(
        source="speechbrain/metricgan-plus-voicebank",
        savedir=SAVEDIR,
        run_opts={"device": "cpu"},
    )
    wrap = MaskNet(m.mods.enhance_model).eval()
    feats = m.compute_features(torch.randn(1, 16000))
    with torch.no_grad():
        torch.onnx.export(
            wrap,
            feats,
            str(out),
            input_names=["feats"],
            output_names=["mask"],
            dynamic_axes={
                "feats": {0: "batch", 1: "time"},
                "mask": {0: "batch", 1: "time"},
            },
            opset_version=17,
            dynamo=False,
        )
    print("wrote", out, "bytes", out.stat().st_size, "example_feats", tuple(feats.shape))


if __name__ == "__main__":
    main()