pi-multifx-pedal/tests/test_nam_host.py

"""Tests for NAM model host — loading, inference, and model switching.

Tests use synthetically generated .nam model files so no
external downloads are required. The test helper `make_test_nam()`
generates valid TorchScript-able NAM models using the nam library.

Uses real NAM model inference via the `nam` Python package,
so these are integration tests, not mocks.
"""

from __future__ import annotations

import json
import tempfile
from pathlib import Path

import numpy as np
import pytest

from src.dsp.nam_host import NAMHost, NAMModel, ModelSwitchMode


# ── Helpers ────────────────────────────────────────────────────────────


def _make_nam_config(arch: str = "Linear", num_weights: int = 1200) -> dict:
    """Create a minimal valid NAM model config dict.

    Generates random weights so the model compiles and runs.
    The resulting .nam file is tiny (~5 KB) and suitable for testing.

    Uses Linear architecture (simplest, supported by init_from_nam).

    Args:
        arch: Model architecture (only "Linear" and "WaveNet" are
              supported by init_from_nam in the nam package).
        num_weights: Number of weight parameters.

    Returns:
        A dict that can be serialized to a .nam file.
    """
    rng = np.random.RandomState(42)

    if arch == "Linear":
        config = {
            "receptive_field": 16,
        }
    elif arch == "WaveNet":
        config = {
            "layers": [{
                "channels": [4] * 3,
                "kernel_size": 3,
                "dilation_base": 2,
                "activation": [{"type": "Tanh"}],
                "gating": True,
                "head": {"channels": [4, 1], "kernel_size": 3},
                "head_scale": 1.0,
            }],
        }
    else:
        raise ValueError(f"Test helper only supports Linear/WaveNet (got {arch})")

    return {
        "version": "0.13.0",
        "architecture": arch,
        "config": config,
        "sample_rate": 48000,
        "weights": rng.uniform(-0.5, 0.5, num_weights).tolist(),
    }


@pytest.fixture
def test_nam_file(tmp_path: Path) -> Path:
    """Create a valid test .nam file and return its path."""
    config = _make_nam_config()
    path = tmp_path / "test_model.nam"
    with open(path, "w") as f:
        json.dump(config, f)
    return path


@pytest.fixture
def host(tmp_path: Path) -> NAMHost:
    """Create a NAMHost with a temp models directory."""
    return NAMHost(models_dir=tmp_path, device="cpu")


# ── Model metadata tests ──────────────────────────────────────────────


class TestNAMModelMetadata:
    def test_minimal_fields(self):
        """NAMModel can be constructed with required fields."""
        model = NAMModel(
            name="test",
            path="/tmp/test.nam",
            architecture="WaveNet",
            size_mb=4.5,
            params_k=6.4,
            receptive_field=64,
            sample_rate=48000,
            compatible=True,
        )
        assert model.name == "test"
        assert model.compatible
        assert model.size_mb == 4.5
        assert model.receptive_field == 64

    def test_large_model_not_compatible(self):
        """Models > 10 MB are flagged as incompatible."""
        model = NAMModel(
            name="big",
            path="/tmp/big.nam",
            architecture="WaveNet",
            size_mb=42.0,
            params_k=500.0,
            receptive_field=512,
            sample_rate=48000,
            compatible=False,
        )
        assert not model.compatible
        assert model.size_mb == 42.0


# ── NAMHost lifecycle tests ───────────────────────────────────────────


class TestNAMHost:
    def test_initial_state(self, host: NAMHost):
        """Fresh host has no model loaded."""
        assert not host.is_loaded
        assert host.current_model is None
        assert host.avg_inference_ms == 0.0

    def test_load_model_success(self, host: NAMHost, test_nam_file: Path):
        """Can load a valid .nam model file."""
        result = host.load_model(str(test_nam_file))
        assert result
        assert host.is_loaded
        assert host.current_model is not None
        assert host.current_model.name == "test_model"
        assert host.current_model.architecture == "Linear"

    def test_load_model_not_found(self, host: NAMHost):
        """Loading a non-existent file returns False."""
        result = host.load_model("/nonexistent/model.nam")
        assert not result
        assert not host.is_loaded

    def test_load_model_bad_extension(self, host: NAMHost, tmp_path: Path):
        """Loading a non-.nam file returns False."""
        bad_file = tmp_path / "model.wav"
        bad_file.write_bytes(b"fake")
        result = host.load_model(str(bad_file))
        assert not result

    def test_unload(self, host: NAMHost, test_nam_file: Path):
        """Unload clears model state."""
        host.load_model(str(test_nam_file))
        assert host.is_loaded

        host.unload()
        assert not host.is_loaded
        assert host.current_model is None

    def test_double_load(self, host: NAMHost, test_nam_file: Path,
                         tmp_path: Path):
        """Loading a second model replaces the first."""
        host.load_model(str(test_nam_file))
        assert host.current_model.name == "test_model"

        # Create second model
        config2 = _make_nam_config()
        path2 = tmp_path / "model2.nam"
        with open(path2, "w") as f:
            json.dump(config2, f)

        host.load_model(str(path2))
        assert host.current_model.name == "model2"
        assert host.is_loaded


# ── Inference tests ───────────────────────────────────────────────────


class TestNAMInference:
    def test_process_1d(self, host: NAMHost, test_nam_file: Path):
        """1D float32 input produces 1D float32 output."""
        host.load_model(str(test_nam_file))
        audio_in = np.random.randn(256).astype(np.float32)
        audio_out = host.process(audio_in)
        assert audio_out.shape == (256,)
        assert audio_out.dtype == np.float32
        # Output should be finite (model may overshoot with random weights)
        assert np.all(np.isfinite(audio_out))

    def test_process_2d(self, host: NAMHost, test_nam_file: Path):
        """2D (1, N) float32 input produces matching 2D output."""
        host.load_model(str(test_nam_file))
        audio_in = np.random.randn(1, 256).astype(np.float32)
        audio_out = host.process(audio_in)
        assert audio_out.shape == (1, 256)
        assert audio_out.dtype == np.float32

    def test_process_no_model(self, host: NAMHost):
        """Processing with no model loaded should pass through."""
        audio_in = np.random.randn(256).astype(np.float32)
        audio_out = host.process(audio_in)
        np.testing.assert_array_equal(audio_out, audio_in)

    def test_process_sine_wave(self, host: NAMHost, test_nam_file: Path):
        """A sine wave should produce a non-silent output."""
        host.load_model(str(test_nam_file))
        t = np.linspace(0, 256 / 48000, 256, dtype=np.float32)
        sine = (np.sin(2 * np.pi * 440 * t) * 0.5).astype(np.float32)
        out = host.process(sine)
        # Should have non-zero energy
        rms_out = np.sqrt(np.mean(out ** 2))
        assert rms_out > 0.0, "Model output should not be silent"

    def test_process_multiple_blocks(self, host: NAMHost,
                                     test_nam_file: Path):
        """Processing multiple blocks should maintain state consistency."""
        host.load_model(str(test_nam_file))
        block = np.random.randn(256).astype(np.float32)

        # Process same block twice
        out1 = host.process(block.copy())
        out2 = host.process(block.copy())

        # Models with no state (ConvNet) should produce same output
        assert out1.shape == out2.shape == (256,)
        assert out1.dtype == np.float32

    def test_process_different_block_sizes(self, host: NAMHost,
                                            test_nam_file: Path):
        """Should handle various block sizes >= receptive field."""
        host.load_model(str(test_nam_file))
        rf = host.current_model.receptive_field

        for block_size in [rf, rf * 2, 128, 256, 512]:
            audio = np.random.randn(block_size).astype(np.float32)
            out = host.process(audio)
            assert out.shape == (block_size,), (
                f"Block size {block_size} produced {out.shape}"
            )


# ── Model switching tests ─────────────────────────────────────────────


class TestModelSwitching:
    def test_instant_switch(self, tmp_path: Path):
        """Instant switch mode should immediately route through new model."""
        host = NAMHost(switch_mode=ModelSwitchMode.INSTANT)

        # Load first model (seed 42 — default)
        c1 = _make_nam_config()
        p1 = tmp_path / "m1.nam"
        with open(p1, "w") as f:
            json.dump(c1, f)
        host.load_model(str(p1))

        audio_in = np.random.randn(256).astype(np.float32)
        out_before = host.process(audio_in)

        # Load second model with DIFFERENT weights by using WaveNet arch
        c2 = _make_nam_config(arch="WaveNet", num_weights=2400)
        p2 = tmp_path / "m2.nam"
        with open(p2, "w") as f:
            json.dump(c2, f)
        host.load_model(str(p2))

        out_after = host.process(audio_in)

        # Different architectures should produce different output
        assert not np.allclose(out_before, out_after), (
            "Different models should produce different output"
        )

    def test_warm_up(self, host: NAMHost, test_nam_file: Path):
        """Warm-up should not crash and load model state."""
        host.load_model(str(test_nam_file))
        # First warm-up after model load
        host.warm_up(256)
        # Should be able to process after warm-up
        out = host.process(np.random.randn(256).astype(np.float32))
        assert out.shape == (256,)

    def test_list_available_models(self, host: NAMHost,
                                   test_nam_file: Path):
        """List available models should find test model."""
        models = host.list_available_models()
        assert len(models) >= 1
        names = [m.name for m in models]
        assert "test_model" in names


# ── Standalone function tests ─────────────────────────────────────────


class TestStandaloneFunctions:
    def test_available_models(self, tmp_path: Path):
        """available_models() returns lightweight model info."""
        from src.dsp.nam_host import available_models

        # Create a test model
        config = _make_nam_config()
        path = tmp_path / "test.nam"
        with open(path, "w") as f:
            json.dump(config, f)

        models = available_models(tmp_path)
        assert len(models) == 1
        assert models[0]["name"] == "test"
        assert models[0]["architecture"] == "Linear"
        assert models[0]["feather"] is True
        assert models[0]["size_mb"] > 0

    def test_available_models_empty_dir(self, tmp_path: Path):
        """Empty directory returns empty list."""
        from src.dsp.nam_host import available_models
        assert available_models(tmp_path) == []

    def test_process_with_model(self, tmp_path: Path):
        """process_with_model convenience function works end-to-end."""
        from src.dsp.nam_host import process_with_model

        config = _make_nam_config()
        path = tmp_path / "test.nam"
        with open(path, "w") as f:
            json.dump(config, f)

        audio_in = np.random.randn(256).astype(np.float32)
        audio_out = process_with_model(str(path), audio_in, device="cpu")
        assert audio_out.shape == audio_in.shape
        assert np.all(np.isfinite(audio_out))


# ── Edge case tests ───────────────────────────────────────────────────


class TestEdgeCases:
    def test_block_smaller_than_rf(self, host: NAMHost,
                                   test_nam_file: Path):
        """Blocks smaller than receptive field are padded."""
        host.load_model(str(test_nam_file))
        rf = host.current_model.receptive_field

        # Create block smaller than receptive field
        small_block = np.random.randn(max(1, rf // 4)).astype(np.float32)
        out = host.process(small_block)
        assert out.shape == small_block.shape
        assert np.all(np.isfinite(out))

    def test_silent_input(self, host: NAMHost, test_nam_file: Path):
        """Silent input should produce valid (possibly non-zero) output."""
        host.load_model(str(test_nam_file))
        silence = np.zeros(256, dtype=np.float32)
        out = host.process(silence)
        assert out.shape == (256,)
        assert np.all(np.isfinite(out))

    def test_model_cache_reuse(self, host: NAMHost, test_nam_file: Path):
        """Loading the same model twice uses cache."""
        path = str(test_nam_file)
        host.load_model(path)
        assert host.is_loaded
        name1 = host.current_model.name

        # Load same path again
        host.load_model(path)
        assert host.current_model.name == name1


# ── Performance tests ─────────────────────────────────────────────────


class TestPerformance:
    def test_inference_timing(self, host: NAMHost, test_nam_file: Path):
        """Track average inference time (should be < 5ms on CPU)."""
        host.load_model(str(test_nam_file))

        # Process many blocks to get stable average
        for _ in range(50):
            block = np.random.randn(256).astype(np.float32)
            host.process(block)

        avg_ms = host.avg_inference_ms
        assert avg_ms > 0.0
        # On modern x86 CPU, ConvNet with 6.4K params should be < 1ms
        # On RPi 4B, expect < 5ms per block (256 samples @ 48kHz = 5.3ms)
        assert avg_ms < 10.0, (
            f"Inference too slow: {avg_ms:.3f}ms (target < 5ms)"
        )

    def test_memory_stability(self, host: NAMHost, test_nam_file: Path):
        """Processing many blocks should not leak or crash."""
        host.load_model(str(test_nam_file))
        for _ in range(200):
            block = np.random.randn(256).astype(np.float32)
            _ = host.process(block)
        # If we get here, no crash
        assert host.avg_inference_ms > 0