"""DeepFilterNet3 推理入口 (AXERA 平台, 当前 AX650N) (依赖仅 numpy + axengine)。

用法:
    python3 inference.py input.wav [-o output.wav] [--model-dir axmodels]
"""
import argparse
import wave
from pathlib import Path

import numpy as np

SR = 48000


def read_wav(path):
    with wave.open(str(path), "rb") as w:
        assert w.getnchannels() == 1, "仅支持单声道"
        sr = w.getframerate()
        data = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16)
    return (data.astype(np.float32) / 32768.0), sr


def write_wav(path, audio, sr):
    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sr)
        w.writeframes(pcm.tobytes())


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path)
    parser.add_argument("-o", "--output", type=Path, default=None)
    parser.add_argument("--model-dir", type=Path, default=Path("axmodels"))
    args = parser.parse_args()

    from deepfilternet3_ax import DeepFilterNet3

    audio, sr = read_wav(args.input)
    if sr != SR:
        raise SystemExit(f"仅支持 48kHz 输入, got {sr}Hz")
    enh = DeepFilterNet3(args.model_dir)
    out = enh.enhance(audio)
    out_path = args.output or args.input.with_name(args.input.stem + "_enhanced.wav")
    write_wav(out_path, out, SR)
    print(f"enhanced -> {out_path}")


if __name__ == "__main__":
    main()