"""Phase 7 — tests de l'encodeur audio from scratch.

CPU, rapides : mel-spectrogramme maison, encodeur, shapes, et entraînement réel
court qui doit transcrire les voyelles synthétiques.
"""

from __future__ import annotations

import numpy as np
import pytest
import torch

from src.data.audio_features import (
    LogMelSpectrogram,
    hz_to_mel,
    mel_filterbank,
    mel_to_hz,
    stft_magnitude,
)
from src.data.prepare_audio import prepare
from src.data.toy_audio import SR, generate_dataset, synth_vowel, to_arrays
from src.model.audio_encoder import AudioConfig, AudioEncoder
from src.tokenizer import load_tokenizer


@pytest.fixture(scope="module")
def tokenizer():
    tok = load_tokenizer()
    if tok.vocab_size != 32000:
        pytest.skip("tokenizer final 32k requis")
    return tok


# ---------------------------------------------------------------------------
# Mel-spectrogramme maison
# ---------------------------------------------------------------------------


def test_hz_mel_roundtrip():
    f = torch.tensor([0.0, 440.0, 1000.0, 8000.0])
    assert torch.allclose(mel_to_hz(hz_to_mel(f)).float(), f, atol=1e-3)


def test_mel_filterbank_shape_and_nonneg():
    fb = mel_filterbank(16000, 256, 64)
    assert fb.shape == (64, 129)
    assert (fb >= 0).all()
    assert fb.sum() > 0


def test_stft_pure_tone_peaks_at_frequency():
    sr, f0 = 16000, 1000.0
    t = torch.arange(4096) / sr
    x = torch.sin(2 * np.pi * f0 * t).unsqueeze(0)
    mag = stft_magnitude(x, n_fft=512, hop_length=128)
    freqs = torch.linspace(0, sr / 2, mag.size(-1))
    peak = freqs[mag[0].mean(0).argmax()]
    assert abs(float(peak) - f0) < 60


def test_logmel_shape():
    wav = torch.from_numpy(to_arrays(generate_dataset(4, seed=0)))
    mel = LogMelSpectrogram(sr=SR, n_fft=256, hop_length=64, n_mels=64)(wav)
    assert mel.shape[0] == 4 and mel.shape[1] == 1 and mel.shape[2] == 64
    assert torch.isfinite(mel).all()


def test_synth_vowel_deterministic_and_bounded():
    a = synth_vowel("a", "grave", rng=np.random.default_rng(0))
    b = synth_vowel("a", "grave", rng=np.random.default_rng(0))
    assert np.allclose(a, b)
    assert np.abs(a).max() <= 1.0 + 1e-6
    assert a.shape == (6144,)


def test_generate_dataset_classes():
    items = generate_dataset(10, seed=0)
    assert len(items) == 10
    assert len({it.caption for it in items}) == 10


# ---------------------------------------------------------------------------
# Données préparées
# ---------------------------------------------------------------------------


def test_prepare_audio_shapes(tmp_path):
    meta = prepare(tmp_path, n_train=20, n_val=10, n_mels=64, n_frames=96)
    assert meta["n_classes"] == 10
    from src.data.prepare_audio import load_mels

    tr, caps = load_mels(tmp_path, "train")
    assert tr.shape == (20, 64, 96)
    assert len(caps) == 20


# ---------------------------------------------------------------------------
# Encodeur
# ---------------------------------------------------------------------------


def test_audio_encoder_shapes():
    enc = AudioEncoder(AudioConfig(n_mels=64, n_frames=96, patch_mels=16,
                                   patch_frames=16, n_layer=1, n_head=2, n_embd=64))
    mel = torch.randn(2, 1, 64, 96)
    tokens = enc(mel)
    # grid 4 x 6 = 24 patchs, pool mean -> pas de CLS
    assert enc.n_audio_tokens == 24
    assert tokens.shape == (2, 24, 64)
    assert enc.pool(tokens).shape == (2, 64)


def test_audio_encoder_is_bidirectional():
    """L'audio n'est pas causal : modifier une trame change les features des autres."""
    torch.manual_seed(0)
    enc = AudioEncoder(AudioConfig(n_mels=64, n_frames=96, patch_mels=16,
                                   patch_frames=16, n_layer=1, n_head=2, n_embd=64)).eval()
    mel = torch.randn(1, 1, 64, 96)
    a = enc(mel)
    mel2 = mel.clone()
    mel2[0, :, :, :16] += 1.0  # modifie la 1re colonne de patchs (trames)
    b = enc(mel2)
    # Les tokens de patchs lointains (dernière colonne) changent aussi -> bidirectionnel.
    assert not torch.allclose(a[:, -1], b[:, -1], atol=1e-4)


# ---------------------------------------------------------------------------
# Entraînement réel court
# ---------------------------------------------------------------------------


def test_audio_model_transcribes(tmp_path, tokenizer):
    from src.train.train_audio import AudioTrainConfig, train

    prepare(tmp_path / "data", n_train=800, n_val=200, n_mels=64, n_frames=96)
    cfg = AudioTrainConfig(
        data_dir=str(tmp_path / "data"), out_dir=str(tmp_path / "ckpt"),
        epochs=400, batch_size=32, lr=3e-3, text_layers=2, text_embd=96,
        text_heads=4, audio_layers=2, audio_embd=96, patch_mels=16,
        patch_frames=16, device="cpu",
    )
    summary = train(cfg)
    assert (tmp_path / "ckpt" / "best.pt").is_file()
    # 10 classes -> hasard 0.10 ; la transcription doit nettement dépasser.
    assert summary["greedy_acc"] > 0.5, summary["examples"]
    assert summary["final_loss"] < 1.5