"""PoC (работа A, этап A0): text-to-audio через AudioLDM2 на CPU.

Главный кандидат плана — Stable Audio Open Small — gated (нужен HF-токен),
поэтому PoC делаем на фолбэке AudioLDM2 (cvssp/audioldm2, CC-BY-NC-SA —
в ОСНОВНОЙ пайплайн с некоммерческими весами не пойдёт, см. docs/plan.md).

Запуск:
  ~/.cache/rpg-ai/venv-pxgpt/bin/python tools/aiaudio/gen_audioldm2.py <id> [seed]

Выход: tools/aiaudio/review/<id>_<seed>.wav (16000 Гц — ресэмпл делает post).
Детерминизм: seed + num_inferences фиксированы, промпт в манифесте.
"""
import sys
from pathlib import Path

import numpy as np
import torch

HERE = Path(__file__).resolve().parent
REVIEW = HERE / "review"

# Промпты PoC: амбиент прудов (туман, вода) против текущего процедурного
# gen.mjs (band-noise ветер + капли, 8 с) — критерий «лучше на слух».
PROMPTS = {
    "amb_ponds": {
        "positive": (
            "field recording, calm misty pond at dusk, gentle water lapping, "
            "soft wind through reeds, occasional distant croaking, eerie "
            "quiet melancholic atmosphere, seamless ambience"
        ),
        "negative": (
            "music, melody, rhythm, drums, speech, voices, sudden loud "
            "noises, harsh noise, distortion"
        ),
    },
    "amb_zvenets": {
        "positive": (
            "field recording, quiet rural settlement at dusk, low distant "
            "metallic bell toll, soft wind, faint rustling, somber "
            "melancholic atmosphere, seamless ambience"
        ),
        "negative": (
            "music, melody, song, rhythm, drums, speech, voices, sirens, "
            "harsh noise, distortion"
        ),
    },
}

DURATION = 10.0  # с — заметно длиннее процедурного лупа (8 с), до crossfade
NUM_INFERENCE_STEPS = 100  # 200 — дефолт качества; на CPU берём 100 для PoC
GUIDANCE = 3.0


def main() -> None:
    if len(sys.argv) < 2 or sys.argv[1] not in PROMPTS:
        print(f"использование: gen_audioldm2.py <{'|'.join(PROMPTS)}> [seed]")
        sys.exit(1)
    name = sys.argv[1]
    seed = int(sys.argv[2]) if len(sys.argv) > 2 else 1
    spec = PROMPTS[name]

    import diffusers  # noqa: F401  (после разбора аргументов — чтобы --help был быстрым)

    torch.manual_seed(seed)
    np.random.seed(seed)

    from diffusers import AudioLDM2Pipeline

    # GPU (RTX 3090) при наличии: fp16, иначе CPU fp32. Seed-детерминизм —
    # по-устройству: fp16-GPU и fp32-CPU дают разные реализации одного seed.
    use_gpu = torch.cuda.is_available()
    dtype = torch.float16 if use_gpu else torch.float32
    pipe = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2", torch_dtype=dtype)
    pipe.to("cuda" if use_gpu else "cpu")
    print(f"устройство: {'cuda fp16' if use_gpu else 'cpu fp32 (12 потоков)'}")
    if not use_gpu:
        torch.set_num_threads(12)

    REVIEW.mkdir(parents=True, exist_ok=True)
    out = REVIEW / f"{name}_{seed}.wav"
    print(f"генерация {name} (seed={seed}, {DURATION} с, шаги {NUM_INFERENCE_STEPS})...")

    audio = pipe(
        spec["positive"],
        negative_prompt=spec["negative"],
        num_inference_steps=NUM_INFERENCE_STEPS,
        guidance_scale=GUIDANCE,
        audio_length_in_s=DURATION,
    ).audios[0]

    # AudioLDM2 отдаёт float [-1..1] 16000 Гц моно; пишем PCM16 без клипа.
    audio = np.clip(audio, -1.0, 1.0)
    pcm = (audio * 32767.0).astype(np.int16)
    import wave

    with wave.open(str(out), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(16000)
        w.writeframes(pcm.tobytes())
    print(f"готово: {out.relative_to(HERE)} ({pcm.shape[0] / 16000:.1f} с, 16 кГц)")


if __name__ == "__main__":
    main()