#!/usr/bin/env python3
"""Dump MetricGAN enhance_model I/O so the OpenVINO spike can convert it."""
from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch
from speechbrain.inference.enhancement import SpectralMaskEnhancement

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


def main() -> None:
    m = SpectralMaskEnhancement.from_hparams(
        source="speechbrain/metricgan-plus-voicebank",
        savedir=SAVEDIR,
        run_opts={"device": "cpu"},
    )
    net = m.mods.enhance_model
    print("enhance_model", type(net).__name__)
    print("device", m.device)
    wav = torch.randn(1, 16000)
    feats = m.compute_features(wav)
    print("wav", tuple(wav.shape), "feats", tuple(feats.shape), feats.dtype)
    lengths = torch.tensor([1.0])
    with torch.no_grad():
        mask = net(feats, lengths=lengths)
    print("mask", tuple(mask.shape), mask.dtype)
    print("modules:")
    for name, mod in net.named_children():
        print(f"  {name}: {type(mod).__name__}")


if __name__ == "__main__":
    main()