pi-multifx-pedal/src/system/defaults.py

"""Default NAM models, IRs, and directory structure for the Pi Multi-FX Pedal.

Provides helpers for:
- Getting canonical default directories for IRs and NAM models
- Generating a basic default IR WAV file
- Generating a minimal pass-through NAM model (Linear, identity)
- Ensuring all default directories and fallback files exist at startup
"""

from __future__ import annotations

import json
import logging
from pathlib import Path

import numpy as np
from scipy.io import wavfile

logger = logging.getLogger(__name__)

# ── Canonical paths ────────────────────────────────────────────────────────────

_DEFAULT_IRS_DIR = Path.home() / ".pedal" / "irs" / "default"
_DEFAULT_MODELS_DIR = Path.home() / ".pedal" / "nam" / "default"
_FACTORY_PRESET_DIR = (
    Path(__file__).resolve().parent.parent.parent / "presets" / "factory"
)


def get_default_irs_dir() -> Path:
    """Return the canonical directory for default IR files.

    This is ``~/.pedal/irs/default/``.
    """
    return _DEFAULT_IRS_DIR


def get_default_models_dir() -> Path:
    """Return the canonical directory for default NAM model files.

    This is ``~/.pedal/nam/default/``.
    """
    return _DEFAULT_MODELS_DIR


# ── Default IR generation ─────────────────────────────────────────────────────

_DEFAULT_IR_SAMPLE_RATE = 48000
_DEFAULT_IR_NUM_SAMPLES = 1024


def generate_default_ir(
    path: Path | None = None,
    sample_rate: int = _DEFAULT_IR_SAMPLE_RATE,
    num_samples: int = _DEFAULT_IR_NUM_SAMPLES,
) -> Path:
    """Generate a simple default impulse response WAV file.

    Produces a short exponentially-decaying impulse (not silence) so
    convolution will colour the sound subtly rather than mute it.

    Parameters
    ----------
    path : Path or None
        Destination path.  Defaults to ``~/.pedal/irs/default/default_ir.wav``.
    sample_rate : int
        Sample rate in Hz (default 48000).
    num_samples : int
        Length of the IR in samples (default 1024, ≈21 ms @ 48 kHz).

    Returns
    -------
    Path
        The path the WAV file was written to.
    """
    if path is None:
        path = _DEFAULT_IRS_DIR / "default_ir.wav"

    path.parent.mkdir(parents=True, exist_ok=True)

    # A short decaying impulse: 1 at sample 0, then exponential decay.
    # Using a mild decay constant so the IR is audibly non-trivial.
    t = np.arange(num_samples, dtype=np.float32)
    ir = np.exp(-t / (num_samples * 0.15))  # decay tail
    ir[0] = 1.0  # initial impulse

    # Normalise to [-1, 1]
    peak = np.max(np.abs(ir))
    if peak > 0:
        ir = ir / peak

    # Write as 16-bit PCM WAV
    wavfile.write(str(path), sample_rate, (ir * 32767).astype(np.int16))

    logger.info("Generated default IR: %s (%d samples, %d Hz)", path, num_samples, sample_rate)
    return path


# ── Pass-through NAM model generation ─────────────────────────────────────────

_NAM_PASSTHROUGH = {
    "architecture": "Linear",
    "config": {
        "receptive_field": 1,
        "bias": False,
    },
    "weights": [1.0],
}


def generate_passthrough_nam(path: Path | None = None) -> Path:
    """Generate a minimal pass-through NAM model file (Linear, identity).

    The model is a ``Conv1d(1, 1, 1)`` with weight = ``[1.0]`` and no bias.
    Loading this model in :class:`~src.dsp.nam_host.NAMHost` results in a
    transparent pass-through — audio is not modified.

    Parameters
    ----------
    path : Path or None
        Destination path.  Defaults to ``~/.pedal/nam/default/pass_through.nam``.

    Returns
    -------
    Path
        The path the .nam file was written to.
    """
    if path is None:
        path = _DEFAULT_MODELS_DIR / "pass_through.nam"

    path.parent.mkdir(parents=True, exist_ok=True)

    with open(path, "w") as f:
        json.dump(_NAM_PASSTHROUGH, f)

    logger.info("Generated pass-through NAM model: %s", path)
    return path


# ── Factory preset directories ────────────────────────────────────────────────

_BASE_DIRS = [
    Path.home() / ".pedal",
    Path.home() / ".pedal" / "presets",
    Path.home() / ".pedal" / "irs",
    Path.home() / ".pedal" / "nam",
]


def _ensure_base_dirs() -> None:
    """Create the base ``~/.pedal`` directory tree if it doesn't exist."""
    for d in _BASE_DIRS:
        d.mkdir(parents=True, exist_ok=True)


# ── Main entry point ──────────────────────────────────────────────────────────


def ensure_defaults_exist() -> bool:
    """Ensure all default directories, IRs, and NAM models exist on disk.

    Creates the following if missing:

    * ``~/.pedal/``, ``~/.pedal/presets/``, ``~/.pedal/irs/``, ``~/.pedal/nam/``
    * ``~/.pedal/irs/default/default_ir.wav`` (exponentially-decaying impulse)
    * ``~/.pedal/nam/default/pass_through.nam`` (Linear identity model)

    Returns
    -------
    bool
        ``True`` if all resources exist or were successfully created.
    """
    try:
        _ensure_base_dirs()

        # Generate default IR if it doesn't exist
        default_ir_path = _DEFAULT_IRS_DIR / "default_ir.wav"
        if not default_ir_path.exists():
            generate_default_ir(default_ir_path)
        else:
            logger.debug("Default IR already exists: %s", default_ir_path)

        # Generate pass-through NAM model if it doesn't exist
        default_nam_path = _DEFAULT_MODELS_DIR / "pass_through.nam"
        if not default_nam_path.exists():
            generate_passthrough_nam(default_nam_path)
        else:
            logger.debug("Default NAM model already exists: %s", default_nam_path)

        return True
    except Exception as exc:
        logger.warning("Failed to ensure defaults exist: %s", exc)
        return False


# ── Standalone CLI ────────────────────────────────────────────────────────────

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    ok = ensure_defaults_exist()
    print(f"Defaults {'OK' if ok else 'FAILED'}")
    print(f"  IRs dir:    {get_default_irs_dir()}")
    print(f"  Models dir: {get_default_models_dir()}")
    if ok:
        print(f"  IR file:    {'EXISTS' if (get_default_irs_dir() / 'default_ir.wav').exists() else 'MISSING'}")
        print(f"  NAM file:   {'EXISTS' if (get_default_models_dir() / 'pass_through.nam').exists() else 'MISSING'}")