"""Phase 7 — préparation des mel-spectrogrammes audio (jouets).

Calcule les log-mel (maison) pour le générateur de voyelles synthétiques et les
met en cache en `.npy` avec les légendes.

Usage :

    python -m src.data.prepare_audio --out-dir data/processed/toy_audio \\
        --n-train 2000 --n-val 200 --n-mels 64 --n-frames 96
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from src.data.audio_features import LogMelSpectrogram
from src.data.toy_audio import SR, generate_dataset, to_arrays


def compute_mels(waveforms: np.ndarray, sr: int, n_fft: int, hop: int, n_mels: int,
                 n_frames: int) -> np.ndarray:
    """Waveforms ``(N, T)`` -> log-mel ``(N, n_mels, n_frames)`` (recadré/paddé)."""
    transform = LogMelSpectrogram(sr=sr, n_fft=n_fft, hop_length=hop, n_mels=n_mels)
    with torch.no_grad():
        mel = transform(torch.from_numpy(waveforms))  # (N, 1, n_mels, T)
    mel = mel.squeeze(1)
    t = mel.size(-1)
    if t >= n_frames:
        mel = mel[:, :, :n_frames]
    else:
        mel = torch.nn.functional.pad(mel, (0, n_frames - t))
    return mel.numpy().astype(np.float32)


def prepare(out_dir: Path, n_train: int, n_val: int, sr: int = SR, n_fft: int = 256,
            hop: int = 64, n_mels: int = 64, n_frames: int = 96, seed: int = 0) -> dict:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    tr = generate_dataset(n_train, seed=seed)
    va = generate_dataset(n_val, seed=seed + 999)
    tr_mel = compute_mels(to_arrays(tr), sr, n_fft, hop, n_mels, n_frames)
    va_mel = compute_mels(to_arrays(va), sr, n_fft, hop, n_mels, n_frames)

    np.save(out_dir / "train_mel.npy", tr_mel)
    np.save(out_dir / "val_mel.npy", va_mel)
    (out_dir / "train_captions.json").write_text(
        json.dumps([it.caption for it in tr], ensure_ascii=False), encoding="utf-8")
    (out_dir / "val_captions.json").write_text(
        json.dumps([it.caption for it in va], ensure_ascii=False), encoding="utf-8")

    classes = sorted({it.caption for it in tr})
    meta = {"source": "toy_vowels", "sr": sr, "n_fft": n_fft, "hop": hop,
            "n_mels": n_mels, "n_frames": n_frames, "n_train": n_train, "n_val": n_val,
            "n_classes": len(classes), "classes": classes}
    (out_dir / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2),
                                       encoding="utf-8")
    return meta


def load_mels(out_dir: Path, split: str = "train"):
    out_dir = Path(out_dir)
    mel = np.load(out_dir / f"{split}_mel.npy")
    caps = json.loads((out_dir / f"{split}_captions.json").read_text(encoding="utf-8"))
    return mel, caps


def main() -> None:
    parser = argparse.ArgumentParser(description="Préparer les mel-spectrogrammes (Phase 7).")
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--n-train", type=int, default=2000)
    parser.add_argument("--n-val", type=int, default=200)
    parser.add_argument("--sr", type=int, default=SR)
    parser.add_argument("--n-fft", type=int, default=256)
    parser.add_argument("--hop", type=int, default=64)
    parser.add_argument("--n-mels", type=int, default=64)
    parser.add_argument("--n-frames", type=int, default=96)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()
    meta = prepare(args.out_dir, args.n_train, args.n_val, args.sr, args.n_fft,
                   args.hop, args.n_mels, args.n_frames, args.seed)
    print(json.dumps({k: v for k, v in meta.items() if k != "classes"},
                     ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()