---

"""Извлечение TitaNet-эмбеддингов скользящим окном по аудиозаписям.

Запуск (в окружении с NeMo):
    python extract_titanet.py /path/to/wav_dir /path/to/out_dir
    python extract_titanet.py /path/to/wav_dir /path/to/out_dir --device cuda --csv

Вход: папка с .wav (любой sr/каналы — приводятся к 16 кГц моно).
Выход на каждую запись:
    <name>.npz : embeddings (M,192) L2-норм., starts_s (M,), ends_s (M,), meta
    <name>.csv : (опц., --csv) строки "<name>_NNNN, v1 v2 ... v192" — как ivectors.csv
"""
from __future__ import annotations

import argparse
from pathlib import Path

import librosa
import numpy as np
import torch
import nemo.collections.asr as nemo_asr

SR = 16000  # TitaNet требует строго 16 кГц моно


def l2_normalize(x: np.ndarray, eps: float = 1e-10) -> np.ndarray:
    return x / np.maximum(np.linalg.norm(x, axis=-1, keepdims=True), eps)


def load_model(device: str):
    model = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(
        "nvidia/speakerverification_en_titanet_large"
    )
    model.eval()
    return model.to(device)


@torch.no_grad()
def window_embeddings(
    wav_path: Path, model, win_s: float, hop_s: float, device: str
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """-> (starts_s (M,), ends_s (M,), embeddings (M,192) L2-норм.)."""
    audio, _ = librosa.load(str(wav_path), sr=SR, mono=True)
    win, hop = int(win_s * SR), int(hop_s * SR)

    if len(audio) < win:  # запись короче окна — одно окно по всей записи
        starts = [0]
    else:
        starts = list(range(0, len(audio) - win + 1, hop))

    starts_s, ends_s, embs = [], [], []
    for start in starts:
        chunk = audio[start:start + win]
        sig = torch.tensor(chunk, dtype=torch.float32, device=device).unsqueeze(0)
        length = torch.tensor([sig.shape[1]], dtype=torch.int64, device=device)
        _, emb = model.forward(input_signal=sig, input_signal_length=length)
        embs.append(emb.squeeze(0).cpu().numpy())
        starts_s.append(start / SR)
        ends_s.append((start + len(chunk)) / SR)

    return (
        np.asarray(starts_s, dtype=np.float64),
        np.asarray(ends_s, dtype=np.float64),
        l2_normalize(np.asarray(embs, dtype=np.float32)),
    )


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("wav_dir", help="папка с .wav")
    ap.add_argument("out_dir", help="куда писать эмбеддинги")
    ap.add_argument("--win", type=float, default=1.5, help="окно, сек")
    ap.add_argument("--hop", type=float, default=0.5, help="шаг, сек")
    ap.add_argument("--device", default="cpu", help="cpu или cuda")
    ap.add_argument("--csv", action="store_true", help="дополнительно писать .csv")
    args = ap.parse_args()

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    wavs = sorted(Path(args.wav_dir).glob("*.wav"))
    if not wavs:
        raise SystemExit(f"в {args.wav_dir} не найдено .wav")
    print(f"найдено записей: {len(wavs)} | device={args.device} "
          f"win={args.win}s hop={args.hop}s")

    model = load_model(args.device)

    for wav in wavs:
        starts_s, ends_s, embs = window_embeddings(
            wav, model, args.win, args.hop, args.device
        )
        np.savez(
            out / f"{wav.stem}.npz",
            embeddings=embs, starts_s=starts_s, ends_s=ends_s,
            meta=dict(model="titanet_large", dim=int(embs.shape[1]),
                      win_s=args.win, hop_s=args.hop, sr=SR),
        )
        if args.csv:
            lines = [
                f"{wav.stem}_{i:04d}," + " ".join(f"{v:.6f}" for v in embs[i])
                for i in range(embs.shape[0])
            ]
            (out / f"{wav.stem}.csv").write_text(
                "filename,i-vector\n" + "\n".join(lines)
            )
        print(f"{wav.name}: {embs.shape[0]} окон × {embs.shape[1]} -> {wav.stem}.npz")

    print(f"\nГотово. Эмбеддинги в {out.resolve()}/")


if __name__ == "__main__":
    main()