#!/usr/bin/env python3
"""WideCodec — self-contained inference (44.1 kHz neural codec, depth-20 decoder).

Encode -> decode any audio to a 44.1 kHz reconstruction. The encoder ingests mono
16 kHz; the frozen FSQ codebook (single codebook, 50 tokens/s, ~0.8 kbps) yields the
code stream; the depth-20 decoder renders it back to 44.1 kHz.

This script + the bundled `neucodec/` package are ALL you need — no other source
repo. Weights (`pytorch_model.bin`) are pulled from the HF model repo on first run.

  pip install torch transformers huggingface_hub local-attention einops librosa soundfile
  huggingface-cli login          # if the repo is private (or export HF_TOKEN=hf_...)
  python infer_widecodec.py --input my.wav --out-dir out
  python infer_widecodec.py --input folder/ --out-dir out    # batch a directory
"""
import argparse
import glob
import os
import sys

# make the bundled `neucodec/` package importable regardless of CWD
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")  # plain HTTPS: Xet can hang on some networks

import librosa
import soundfile as sf
import torch

from neucodec import NeuCodec

AUDIO_EXTS = (".wav", ".mp3", ".flac", ".m4a", ".ogg", ".opus")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, help="audio file OR a directory of audio")
    ap.add_argument("--out-dir", default="out")
    ap.add_argument("--repo", default="Scicom-intl/WideCodec")
    ap.add_argument("--depth", type=int, default=20, help="decoder depth (WideCodec=20)")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    a = ap.parse_args()

    os.makedirs(a.out_dir, exist_ok=True)
    print(f"[infer] loading {a.repo} (decoder_depth={a.depth}) on {a.device} …")
    model = NeuCodec._from_pretrained(
        model_id=a.repo, decoder_depth=a.depth, token=os.environ.get("HF_TOKEN")
    ).eval().to(a.device)
    sr_out = model.sample_rate
    print(f"[infer] output sample_rate = {sr_out}")

    if os.path.isdir(a.input):
        files = sorted(f for f in glob.glob(os.path.join(a.input, "*"))
                       if f.lower().endswith(AUDIO_EXTS))
    else:
        files = [a.input]
    print(f"[infer] {len(files)} file(s)")

    for f in files:
        wav16, _ = librosa.load(f, sr=16000, mono=True)  # frozen encoder wants 16 kHz mono
        x = torch.from_numpy(wav16).float().view(1, 1, -1).to(a.device)
        with torch.no_grad():
            codes = model.encode_code(x)
            wav = model.decode_code(codes).squeeze().detach().cpu().float().numpy()
        base = os.path.splitext(os.path.basename(f))[0]
        out = os.path.join(a.out_dir, f"{base}_44k.wav")
        sf.write(out, wav, sr_out)
        print(f"[infer] {f} -> {out}  ({len(wav) / sr_out:.2f}s @ {sr_out} Hz)")


if __name__ == "__main__":
    main()