tests/test_inversion_v2.py

"""CPU-only mechanical contract for the frozen v2 inversion policy.

Expected public API: music3lab.inversion_v2 exposes the pure functions/classes
called below. Metric inputs are [batch, 2, samples] float tensors. The v1.3
loader returns all four bound E2 final latents, never a selected-only latent.
"""
from __future__ import annotations

import hashlib
from pathlib import Path

import pytest
import torch
from safetensors.torch import save as save_safetensors

from music3lab import inversion_v2 as v2


def _audio(samples: int = 128) -> torch.Tensor:
    time = torch.arange(samples, dtype=torch.float32) / samples
    return torch.stack((0.25 * torch.sin(2 * torch.pi * 7 * time), 0.20 * torch.cos(2 * torch.pi * 5 * time))).unsqueeze(0)


def test_time_nmse_and_stft_distinguish_identity_distortion_and_polarity() -> None:
    target = _audio()
    assert torch.equal(v2.target_energy_normalized_time_nmse(target, target), torch.zeros(1))
    assert v2.target_energy_normalized_time_nmse(target + 0.2, target).item() > 0.1
    assert v2.mrstft_magnitude_distance(-target, target, fft_sizes=(32, 64)).item() < 1e-6
    assert v2.time_and_complex_stft_distance(-target, target, fft_sizes=(32, 64)).item() > 0.5


def test_selection_score_is_identical_at_every_training_stage() -> None:
    metrics = {
        "time_nmse": 0.125,
        "complex_stft_nmse": 0.25,
        "legacy_mrstft": 0.375,
        "mid_side_nmse": 0.2,
        "relative_envelope": 0.1,
    }
    scores = [v2.selection_score(metrics, stage=stage) for stage in ("fp32", "bf16", "handoff")]
    assert scores == [scores[0]] * 3


@pytest.mark.parametrize(("stage", "warmup", "total", "peak"), (("fp32", 4, 20, 0.02), ("bf16", 2, 10, 0.01)))
def test_warmup_cosine_schedule_has_exact_frozen_endpoints(stage: str, warmup: int, total: int, peak: float) -> None:
    minimum = peak / 10
    schedule = v2.FrozenStageSchedule(
        name=stage,
        warmup_steps=warmup,
        total_steps=total,
        maximum_learning_rate=peak,
        minimum_learning_rate=minimum,
    )
    assert v2.warmup_cosine_learning_rate(schedule, 0) == 0.0
    assert v2.warmup_cosine_learning_rate(schedule, warmup) == peak
    with pytest.raises(ValueError, match="step"):
        v2.warmup_cosine_learning_rate(schedule, total + 1)


def _authority(path: Path) -> dict[str, object]:
    data = path.read_bytes()
    return {"path": path.name, "sha256": hashlib.sha256(data).hexdigest(), "size": len(data), "required_restart_count": 4, "tensor_key": "final_latents"}


def test_v11_e2_loader_requires_all_four_bound_final_latents(tmp_path: Path) -> None:
    path = tmp_path / "latents.safetensors"
    expected = torch.arange(24, dtype=torch.float32).reshape(4, 2, 3)
    path.write_bytes(save_safetensors({"final_latents": expected, "selected_latent": expected[:1].clone()}))
    authority = _authority(path)
    path.chmod(0o644)
    assert torch.equal(v2.load_v11_e2_final_latents(path, authority), expected)
    selected_only = tmp_path / "selected-only.safetensors"
    selected_only.write_bytes(save_safetensors({"selected_latent": expected[:1]}))
    selected_only.chmod(0o644)
    with pytest.raises((ValueError, RuntimeError), match="final|four|restart|sha|authority"):
        v2.load_v11_e2_final_latents(selected_only, authority)
    cherry_picked = tmp_path / "cherry-picked.safetensors"
    cherry_picked.write_bytes(save_safetensors({"final_latents": expected[:3]}))
    cherry_picked.chmod(0o644)
    with pytest.raises((ValueError, RuntimeError), match="four|restart|sha"):
        v2.load_v11_e2_final_latents(cherry_picked, authority)
    path.write_bytes(save_safetensors({"final_latents": expected + 1}))
    with pytest.raises((ValueError, RuntimeError), match="sha|authority|size"):
        v2.load_v11_e2_final_latents(path, authority)


def test_restart_ledger_is_monotone_and_breaks_objective_ties_by_index() -> None:
    ledger = v2.RestartLedger(restart_count=4)
    ledger.record((4.0, 3.0, 2.0, 1.0))
    ledger.record((5.0, 2.0, 3.0, 1.0))
    assert ledger.best_so_far == (4.0, 2.0, 2.0, 1.0)
    assert ledger.select((1.0, 1.0, 2.0, 3.0)) == 0
    with pytest.raises((ValueError, RuntimeError), match="regress|monotone"):
        ledger.record_best_so_far((4.1, 2.0, 2.0, 1.0))


def test_evaluator_cannot_run_before_lock_or_change_locked_selection() -> None:
    lock = v2.SelectionLock()
    with pytest.raises((ValueError, RuntimeError), match="lock|select"):
        lock.evaluate(lambda index: {"rank": index})
    assert lock.lock((3.0, 1.0, 1.0, 2.0)) == 1
    assert lock.evaluate(lambda index: {"rank": index}) == {"rank": 1}
    with pytest.raises((ValueError, RuntimeError), match="locked|selection"):
        lock.lock((0.0, 3.0, 3.0, 3.0))


def test_distributional_prior_has_no_oracle_distance_input() -> None:
    latents = torch.tensor([[[0.0, 1.0, 9.0, -9.0]]], dtype=torch.float32)
    prior = v2.distributional_prior(latents, mean=0.0, std=1.0, tail_threshold=3.0)
    assert prior["mean_error"] == pytest.approx(0.25)
    assert prior["mean_term"] == pytest.approx(0.25**2)
    assert prior["tail_term"] == pytest.approx(18.0)
    expected = prior["mean_term"] + prior["std_term"] + 0.1 * prior["tail_term"]
    assert prior["prior"] == pytest.approx(expected)
    assert "oracle" not in " ".join(prior.keys()).lower()


class _TinyFrozenDecoder(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.projection = torch.nn.Conv1d(2, 2, 1, bias=False)
        torch.nn.init.eye_(self.projection.weight[..., 0])
        for parameter in self.parameters():
            parameter.requires_grad_(False)

    def forward(self, latents: torch.Tensor) -> torch.Tensor:
        return self.projection(latents)


def test_frozen_decoder_keeps_weights_grad_free_but_latents_receive_gradient() -> None:
    decoder = _TinyFrozenDecoder()
    latents = torch.full((1, 2, 32), 0.1, requires_grad=True)
    audit = v2.frozen_decoder_step(decoder, latents, torch.zeros_like(latents))
    assert audit.weight_sha256_before == audit.weight_sha256_after
    assert audit.finite_loss and audit.finite_latent_gradient and audit.nonzero_latent_gradient
    assert all(parameter.grad is None for parameter in decoder.parameters())


def test_progress_never_substitutes_for_feasibility_or_high_fidelity() -> None:
    failed = v2.quality_gate(progress_fraction=0.99, feasibility_metrics={"time_nmse": 2.0}, high_fidelity_metrics={"si_sdr_db": -10.0})
    assert not failed.feasibility_pass
    assert not failed.high_fidelity_pass
    assert failed.quality_claim == "NO_QUALITY_CLAIM"


def test_authority_tamper_and_atomic_publication(tmp_path: Path) -> None:
    staged = tmp_path / "staged"
    staged.mkdir()
    staged.chmod(0o755)
    payload = b"v2 evidence"
    (staged / "metrics.json").write_bytes(payload)
    (staged / "metrics.json").chmod(0o644)
    authority = {"metrics.json": hashlib.sha256(payload).hexdigest()}
    published = tmp_path / "published"
    v2.atomic_publish_verified(staged, published, authority)
    assert (published / "metrics.json").read_bytes() == payload
    assert not staged.exists()
    (published / "metrics.json").write_bytes(b"tampered")
    with pytest.raises((ValueError, RuntimeError), match="sha|authority|tamper"):
        v2.verify_published_authority(published, authority)
    replacement = tmp_path / "replacement"
    replacement.mkdir()
    replacement.chmod(0o755)
    (replacement / "metrics.json").write_bytes(payload)
    (replacement / "metrics.json").chmod(0o644)
    with pytest.raises((ValueError, RuntimeError), match="authority|sha"):
        v2.atomic_publish_verified(replacement, published, {"metrics.json": "0" * 64})
    assert (published / "metrics.json").read_bytes() == b"tampered"