"""NAM A2 model host — load, configure, and run inference on RPi 4B.

Leverages the neural-amp-modeler (NAM) .nam file format for model loading
and inference via a lightweight PyTorch reimplementation. Supports:
  - SlimmableContainer (A2 format): version 0.12+ with quality tier submodels
  - WaveNet: standard dilated-convolution amp model
  - Linear: simple pass-through/bias model

On RPi 4B:
- Python + PyTorch: good for Feather/Nano models only (~0.5-3ms at 256-block)
- For Standard/Lite models, use neural-amp-modeler-lv2 compiled natively
  (NeuralAudio engine, LV2 plugin, ~1-2ms at 256-block)
- For A2 Slimmable runtime quality dialing, port to OpenSauce/nam-rs
"""

from __future__ import annotations

import json
import logging
import time
import warnings
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Optional

import numpy as np


# ── Architecture-aware NAM model builder ──────────────────────────────

def _build_nam_model(data: dict) -> "torch.nn.Module":
    """Build a PyTorch model from a parsed .nam file dict.

    Dispatches to the appropriate builder based on ``architecture``:

    * ``SlimmableContainer`` — A2 format: pick the highest-quality
      submodel (``max_value`` closest to 1.0) and build its WaveNet.
    * ``WaveNet`` — standard dilated-convolution architecture.
    * ``Linear`` — simple single-weight pass-through.

    Args:
        data: Parsed .nam JSON content.

    Returns:
        A ``torch.nn.Module`` ready for inference.
    """
    import torch.nn as nn

    arch = data.get("architecture", "Linear")
    config = data.get("config", {})
    weights = data.get("weights", [])

    if arch == "SlimmableContainer":
        # Pick the highest-quality submodel
        submodels = config.get("submodels", [])
        if not submodels:
            raise ValueError("SlimmableContainer has no submodels")
        # Sort by max_value descending — pick highest quality
        best = max(submodels, key=lambda sm: sm.get("max_value", 0))
        model_data = best.get("model", {})
        if not model_data:
            raise ValueError("SlimmableContainer submodel has no 'model' data")
        # Recurse — the submodel's model field is a complete .nam export dict
        return _build_nam_model(model_data)

    elif arch == "WaveNet":
        return _build_wavenet(config, weights)

    elif arch == "Linear":
        return _build_linear(config, weights)

    else:
        raise ValueError(f"Unsupported NAM architecture: {arch!r}")


def _build_linear(config: dict, weights: list) -> "torch.nn.Module":
    """Build a Linear pass-through model.

    NAM Linear model: just a scalar gain and optional bias.
    """
    import torch
    import torch.nn as nn

    has_bias = config.get("bias", False)

    class LinearNAM(nn.Module):
        def __init__(self, gain: float, bias: float = 0.0):
            super().__init__()
            self.gain = nn.Parameter(torch.tensor(gain, dtype=torch.float32))
            self.bias = nn.Parameter(torch.tensor(bias, dtype=torch.float32)) if has_bias else 0.0

        def forward(self, x):
            # x: (B, T) or (T,)
            if isinstance(self.bias, nn.Parameter):
                return x * self.gain + self.bias
            return x * self.gain

    gain = float(weights[0]) if weights else 1.0
    bias = float(weights[1]) if has_bias and len(weights) > 1 else 0.0
    return LinearNAM(gain, bias)


def _build_wavenet(config: dict, weights: list) -> "torch.nn.Module":
    """Build a WaveNet NAM model from config and flat weight array.

    Implements the standard NAM WaveNet architecture:
    - RechannelIn: 1x1 conv (1 → C)
    - N dilated conv layers with residual connections
    - HeadRechannel: temporal conv (C → 1) + optional bias
    - Optional WaveNet-level Head (activation + conv blocks)
    - head_scale multiplier

    Weights are imported in the exact order matching NAM's export format:

    1. rechannel (weight only, no bias)
    2. Per layer: conv(w+b), input_mixer(w, no bias)
       (layer1x1, head1x1, FiLM omitted — not present in A2 models)
    3. head_rechannel (w+b)
    4. Optional head (conv w+b per block)
    5. head_scale (single float)
    """
    import torch
    import torch.nn as nn

    layers_cfg = config.get("layers", [])
    head_cfg = config.get("head", None)  # WaveNet-level head
    head_scale = config.get("head_scale", 1.0)

    if not layers_cfg:
        raise ValueError("WaveNet config has no layers")

    # ── Extract layer array config (we use first/only layer array) ─────
    la = layers_cfg[0]  # A2 models have single layer array
    input_size = la.get("input_size", 1)
    condition_size = la.get("condition_size", 1)
    channels = la.get("channels", 8)
    kernel_sizes = la.get("kernel_sizes", la.get("kernel_size", [3]))
    dilations = la.get("dilations", [1])
    la_head_cfg = la.get("head", {"out_channels": 1, "kernel_size": 1, "bias": True})

    # Normalize scalar kernel_size to list
    if isinstance(kernel_sizes, int):
        kernel_sizes = [kernel_sizes] * len(dilations)
    if isinstance(dilations, int):
        dilations = [dilations] * len(kernel_sizes)

    num_layers = len(dilations)
    out_channels = la_head_cfg.get("out_channels", 1)
    head_kernel_size = la_head_cfg.get("kernel_size", 1)
    head_bias = la_head_cfg.get("bias", True)

    # ── Build modules ──────────────────────────────────────────────────

    modules = []

    # 1. RechannelIn: 1x1 conv, no bias
    rechannel = nn.Conv1d(input_size, channels, 1, bias=False)
    modules.append(("rechannel", rechannel))

    # 2. Dilated conv layers + optional layer1x1
    layer_convs = nn.ModuleList()
    input_mixers = nn.ModuleList()
    layer1x1s = nn.ModuleList()
    for i in range(num_layers):
        ks = kernel_sizes[i] if i < len(kernel_sizes) else kernel_sizes[-1]
        d = dilations[i] if i < len(dilations) else 1
        conv = nn.Conv1d(channels, channels, ks, dilation=d, padding=0)
        im = nn.Conv1d(condition_size, channels, 1, bias=False)
        layer_convs.append(conv)
        input_mixers.append(im)
        # Many A2 models include a 1x1 post-conv (layer1x1)
        l1 = nn.Conv1d(channels, channels, 1, bias=True)
        layer1x1s.append(l1)

    modules.append(("layer_convs", layer_convs))
    modules.append(("input_mixers", input_mixers))
    modules.append(("layer1x1s", layer1x1s))

    # 3. HeadRechannel: temporal conv
    head_rechannel = nn.Conv1d(channels, out_channels, head_kernel_size, bias=head_bias)
    modules.append(("head_rechannel", head_rechannel))

    # 4. Optional WaveNet-level head
    head_module = None
    if head_cfg is not None and head_cfg.get("kernel_size", 0) > 0:
        h_ks = head_cfg.get("kernel_size", 1)
        h_bias = head_cfg.get("bias", True)
        head_module = nn.Conv1d(out_channels, out_channels, h_ks, bias=h_bias)
        modules.append(("head", head_module))

    # ── Assemble WaveNetNAM module ─────────────────────────────────────
    class WaveNetNAM(nn.Module):
        def __init__(self, rechannel, layer_convs, input_mixers, layer1x1s,
                     head_rechannel, head, head_scale, receptive_field):
            super().__init__()
            self.rechannel = rechannel
            self.layer_convs = layer_convs
            self.input_mixers = input_mixers
            self.layer1x1s = layer1x1s
            self.head_rechannel = head_rechannel
            self.head = head
            self.head_scale = head_scale
            self.receptive_field = receptive_field

        def forward(self, x):
            # x: (B, T) or (B, 1, T)
            if x.dim() == 2:
                x = x.unsqueeze(1)  # (B, 1, T)

            # Rechannel
            x = self.rechannel(x)  # (B, C, T)

            # Dilated conv layers with residual
            for conv, im, l1 in zip(self.layer_convs, self.input_mixers, self.layer1x1s):
                # Causal padding (left-only) for dilated conv
                pad = (conv.kernel_size[0] - 1) * conv.dilation[0]
                if pad > 0:
                    x_pad = nn.functional.pad(x, (pad, 0))
                else:
                    x_pad = x
                # Dilated conv
                out = conv(x_pad)
                # Input mixer (condition path)
                cond = torch.ones_like(x[:, :1, :])
                c = im(cond)
                # Ensure same length
                if out.shape[-1] < x.shape[-1]:
                    x = x[:, :, -out.shape[-1]:]
                if c.shape[-1] < x.shape[-1]:
                    c = c[:, :, -x.shape[-1]:]
                # Gated activation
                out = torch.tanh(out + c)
                # 1x1 post-conv (layer1x1)
                out = l1(out)
                # Truncate to out length if needed
                if out.shape[-1] < x.shape[-1]:
                    x = x[:, :, -out.shape[-1]:]
                # Residual
                x = x + out

            # Head rechannel
            x = self.head_rechannel(x)  # (B, 1, T')
            if x.shape[-1] > 0:
                # Trim to receptive field boundary
                pass

            # Optional head
            if self.head is not None:
                pad = (self.head.kernel_size[0] - 1)
                if pad > 0:
                    x = nn.functional.pad(x, (pad, 0))
                x = self.head(x)

            # Scale
            x = x * self.head_scale

            return x.squeeze(1)  # (B, T')

    # Compute receptive field
    rf = 1
    for i in range(num_layers):
        ks = kernel_sizes[i] if i < len(kernel_sizes) else kernel_sizes[-1]
        d = dilations[i] if i < len(dilations) else 1
        rf += (ks - 1) * d
    rf += head_kernel_size - 1
    if head_module is not None:
        rf += head_cfg.get("kernel_size", 1) - 1

    model = WaveNetNAM(
        rechannel, layer_convs, input_mixers, layer1x1s,
        head_rechannel, head_module, head_scale, rf,
    )

    # ── Import weights from flat array ─────────────────────────────────
    _import_wavenet_weights(model, weights)
    return model


def _import_wavenet_weights(model: "torch.nn.Module", weights: list) -> None:
    """Import flat weight array into a WaveNetNAM model.

    Weight order (matching NAM's export_weights):
    1. rechannel.weight
    2. For each layer: conv.weight, conv.bias, input_mixer.weight
    3. head_rechannel.weight, head_rechannel.bias
    4. head.weight, head.bias (if head exists)
    5. head_scale
    """
    import torch as _torch
    import numpy as _np

    w = _torch.tensor(weights, dtype=_torch.float32)
    i = 0

    # rechannel.weight: [C, 1, 1]
    n = model.rechannel.weight.numel()
    model.rechannel.weight.data = w[i:i+n].reshape(model.rechannel.weight.shape)
    i += n

    # Per-layer weights
    for conv, im, l1 in zip(model.layer_convs, model.input_mixers, model.layer1x1s):
        # conv.weight: [C, C, ks]
        n = conv.weight.numel()
        conv.weight.data = w[i:i+n].reshape(conv.weight.shape)
        i += n
        # conv.bias: [C]
        n = conv.bias.numel()
        conv.bias.data = w[i:i+n].reshape(conv.bias.shape)
        i += n
        # input_mixer.weight: [C, condition_size, 1]
        n = im.weight.numel()
        im.weight.data = w[i:i+n].reshape(im.weight.shape)
        i += n
        # layer1x1.weight: [C, C, 1]
        n = l1.weight.numel()
        l1.weight.data = w[i:i+n].reshape(l1.weight.shape)
        i += n
        # layer1x1.bias: [C]
        n = l1.bias.numel()
        l1.bias.data = w[i:i+n].reshape(l1.bias.shape)
        i += n

    # head_rechannel.weight
    n = model.head_rechannel.weight.numel()
    model.head_rechannel.weight.data = w[i:i+n].reshape(model.head_rechannel.weight.shape)
    i += n
    # head_rechannel.bias
    if model.head_rechannel.bias is not None:
        n = model.head_rechannel.bias.numel()
        model.head_rechannel.bias.data = w[i:i+n].reshape(model.head_rechannel.bias.shape)
        i += n

    # head weight + bias (if present)
    if model.head is not None:
        n = model.head.weight.numel()
        model.head.weight.data = w[i:i+n].reshape(model.head.weight.shape)
        i += n
        if model.head.bias is not None:
            n = model.head.bias.numel()
            model.head.bias.data = w[i:i+n].reshape(model.head.bias.shape)
            i += n

    # head_scale (final float)
    if i < len(w):
        model.head_scale = float(w[i])
        i += 1

    # Verify we consumed all weights
    if i != len(w):
        warnings.warn(
            f"NAM weight import consumed {i}/{len(w)} weights. "
            f"Model has {sum(p.numel() for p in model.parameters())} params. "
            f"Check weight order matches architecture."
        )


logger = logging.getLogger(__name__)

DEFAULT_NAM_DIR = Path.home() / ".pedal" / "nam"

# ── Enum ────────────────────────────────────────────────────────────────


class ModelSwitchMode(Enum):
    """How to handle audio when switching between NAM models.

    INSTANT:  Immediate switch — may produce a click/pop.
    CROSSFADE: Smooth 256-sample fade between old and new output.
    PAUSE:    Brief silence (~1 block) during model load.
    """
    INSTANT = "instant"
    CROSSFADE = "crossfade"
    PAUSE = "pause"

    def __str__(self) -> str:
        return self.value


# ── Metadata ─────────────────────────────────────────────────────────────


@dataclass
class NAMModel:
    """Metadata for a loaded NAM model."""
    name: str
    path: str
    size_mb: float
    architecture: str
    params_k: float = 0.0
    receptive_field: int = 0
    sample_rate: int = 48000
    compatible: bool = True

    @property
    def family(self) -> str:
        """Categorise the model by file size."""
        if self.size_mb < 0.1:
            return "nano"
        elif self.size_mb < 1.0:
            return "feather"
        elif self.size_mb < 4.0:
            return "lite"
        elif self.size_mb < 10.0:
            return "standard"
        else:
            return "heavy"

    @property
    def estimated_latency_ms(self) -> str:
        """Return estimated per-block latency on RPi 4B at 256-block / 48 kHz."""
        estimates = {
            "nano": "0.1-0.2 ms (always safe)",
            "feather": "0.5-1 ms (safe)",
            "lite": "1-2 ms (OK with compiled, marginal with Python)",
            "standard": "2-4 ms (compiled only)",
            "heavy": "5-10 ms (too expensive for RPi 4B)",
        }
        return estimates.get(self.family, "unknown")


# ── Host ─────────────────────────────────────────────────────────────────


class NAMHost:
    """Hosts NAM models for real-time amp simulation.

    Loads ``.nam`` files (JSON format with weights) and runs inference
    through PyTorch.  On RPi 4B with Python, limit to Feather/Nano
    models for reliable < 10 ms block processing.

    Parameters
    ----------
    models_dir : str | Path
        Directory scanned for available models.
    device : str
        Torch device string (``'cpu'``, ``'cuda'``, …).
    switch_mode : ModelSwitchMode
        Transition style when switching models.
    """

    def __init__(
        self,
        models_dir: str | Path = DEFAULT_NAM_DIR,
        device: str = "cpu",
        switch_mode: ModelSwitchMode = ModelSwitchMode.CROSSFADE,
    ) -> None:
        self._models_dir = Path(models_dir)
        self._device = device
        self._switch_mode = switch_mode

        self._loaded_model: Optional[NAMModel] = None
        self._inference_model = None       # nn.Module instance
        self._torch = None                 # lazy import
        self._torch_device = None          # resolved device object

        # Timing stats
        self._timing_samples: list[float] = []

        # Crossfade state
        self._crossfade_len: int = 256
        self._crossfade_buf: Optional[np.ndarray] = None
        self._crossfade_pos: int = 0

        # Simple model cache (path -> model instance)
        self._model_cache: dict[str, object] = {}

        self._models_dir.mkdir(parents=True, exist_ok=True)

        logger.info(
            "NAMHost(models_dir=%s, device=%s, switch_mode=%s)",
            self._models_dir, device, switch_mode,
        )

    # ── Lazy torch import ──────────────────────────────────────────

    def _import_torch(self):
        if self._torch is None:
            import torch
            self._torch = torch
            self._torch_device = torch.device(self._device)

    # ── Properties ─────────────────────────────────────────────────

    @property
    def is_loaded(self) -> bool:
        return self._loaded_model is not None

    @property
    def current_model(self) -> Optional[NAMModel]:
        return self._loaded_model

    @property
    def avg_inference_ms(self) -> float:
        """Rolling average inference time in ms (last 50 blocks)."""
        if not self._timing_samples:
            return 0.0
        return float(np.mean(self._timing_samples[-50:]))

    # ── Model loading ──────────────────────────────────────────────

    def load_model(self, model_path: str) -> bool:
        """Load a ``.nam`` model file and build its inference model.

        Returns ``True`` on success, ``False`` on error.
        """
        path = Path(model_path)
        if not path.exists() or path.suffix.lower() not in (".nam",):
            logger.error("Model not found or invalid: %s", model_path)
            return False

        # Read file
        try:
            with open(path, "r") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError) as e:
            logger.error("Failed to parse .nam file: %s", e)
            return False

        architecture = data.get("architecture", "Linear")
        config = data.get("config", {})
        size_mb = path.stat().st_size / (1024 * 1024)

        # Build metadata
        self._loaded_model = NAMModel(
            name=path.stem,
            path=str(path),
            size_mb=size_mb,
            architecture=architecture,
            receptive_field=config.get("receptive_field", 0),
        )

        # Build PyTorch inference model
        model_ok = self._build_inference(data)

        if model_ok:
            # Store param count in metadata
            params = sum(p.numel() for p in self._inference_model.parameters())
            self._loaded_model.params_k = round(params / 1000, 1)

        logger.info(
            "Loaded NAM model: %s (%.1f MB, %s, %s family, rf=%d, params=%.1fK)",
            self._loaded_model.name,
            size_mb,
            architecture,
            self._loaded_model.family,
            self._loaded_model.receptive_field,
            self._loaded_model.params_k,
        )
        return True

    def _build_inference(self, data: dict) -> bool:
        """Instantiate a PyTorch model from a NAM config dict.

        Uses our lightweight in-process builder that handles
        SlimmableContainer (A2), WaveNet, and Linear architectures.
        """
        self._import_torch()

        path = data.get("path", "")
        cache_key = str(path)

        # Check cache first
        if cache_key and cache_key in self._model_cache:
            self._inference_model = self._model_cache[cache_key]
            self._inference_model.eval()
            return True

        try:
            model = _build_nam_model(data)
            model.eval()

            # Move to target device
            if str(self._torch_device) != "cpu":
                model = model.to(self._torch_device)

            # Cache it
            if cache_key:
                self._model_cache[cache_key] = model

            self._inference_model = model
            return True

        except ImportError as exc:
            logger.warning("Required package not installed; inference unavailable: %s", exc)
            self._inference_model = None
            return False
        except Exception as exc:
            logger.warning("Failed to build inference model: %s", exc)
            self._inference_model = None
            return False

    def unload(self) -> None:
        """Unload the current model and free resources."""
        self._loaded_model = None
        if self._inference_model is not None:
            self._inference_model = self._inference_model.to("cpu")
            del self._inference_model
            self._inference_model = None
        self._torch = None
        self._torch_device = None
        self._timing_samples.clear()
        self._crossfade_buf = None
        self._crossfade_pos = 0
        logger.info("NAM model unloaded")

    # ── Warm-up ────────────────────────────────────────────────────

    def warm_up(self, block_size: int = 256) -> None:
        """Run a dry inference to warm caches / GPU.

        Safe to call even when no model is loaded.
        """
        if self._inference_model is None:
            return
        self._import_torch()
        try:
            dummy = self._torch.randn(1, block_size, device=self._torch_device)
            with self._torch.no_grad():
                self._inference_model(dummy)
            logger.debug("Model warm-up complete (block=%d)", block_size)
        except Exception as exc:
            logger.warning("Warm-up failed (non-fatal): %s", exc)

    # ── Inference ──────────────────────────────────────────────────

    def process(self, audio_block: np.ndarray) -> np.ndarray:
        """Run a block of audio through the loaded NAM model.

        Handles 1-D (``(N,)``) and 2-D (``(1, N)``) float32 input.
        If no model is loaded, passes audio through unchanged.

        Parameters
        ----------
        audio_block : np.ndarray
            PCM samples, float32 in [-1, 1].

        Returns
        -------
        np.ndarray
            Processed audio, same shape.
        """
        if self._inference_model is None:
            # Pass-through when no model is loaded
            return audio_block

        self._import_torch()
        start = time.perf_counter()

        with self._torch.no_grad():
            x = self._torch.from_numpy(audio_block.astype(np.float32))

            # Ensure shape (1, T) for ConvNet-like models
            was_1d = x.dim() == 1
            if was_1d:
                x = x.unsqueeze(0)

            # Move to device
            if str(self._torch_device) != "cpu":
                x = x.to(self._torch_device)

            y = self._inference_model(x)

            # Squeeze back
            if was_1d:
                y = y.squeeze(0)

            # Move back to CPU numpy
            if str(self._torch_device) != "cpu":
                y = y.cpu()
            y = y.numpy()

        # Trim / pad to match input length
        if y.shape[-1] > audio_block.shape[-1]:
            y = y[..., :audio_block.shape[-1]]
        elif y.shape[-1] < audio_block.shape[-1]:
            pad_len = audio_block.shape[-1] - y.shape[-1]
            y = np.pad(y, ((0, 0),) * (y.ndim - 1) + ((0, pad_len),), mode="constant")

        # Track timing
        elapsed_ms = (time.perf_counter() - start) * 1000
        self._timing_samples.append(elapsed_ms)
        if len(self._timing_samples) > 200:
            self._timing_samples = self._timing_samples[-100:]

        return y.astype(np.float32)

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

    def switch_model(
        self,
        model_path: str,
        mode: Optional[ModelSwitchMode] = None,
    ) -> bool:
        """Load a new model with a configurable transition.

        Parameters
        ----------
        model_path : str
            Path to the new ``.nam`` file.
        mode : ModelSwitchMode or None
            Override the default switch mode for this call.

        Returns
        -------
        bool
            ``True`` on success.
        """
        effective_mode = mode if mode is not None else self._switch_mode

        if effective_mode == ModelSwitchMode.PAUSE:
            # Replace immediately (caller is expected to mute output briefly)
            return self.load_model(model_path)

        if effective_mode == ModelSwitchMode.CROSSFADE and self._inference_model is not None:
            # Snapshot the current output buffer for crossfade
            if self._crossfade_buf is None:
                self._crossfade_buf = np.zeros(self._crossfade_len, dtype=np.float32)
            self._crossfade_pos = 0

        return self.load_model(model_path)

    # ── Model discovery ─────────────────────────────────────────────

    def list_available_models(self) -> list[NAMModel]:
        """Scan ``models_dir`` for ``.nam`` files and return metadata."""
        models: list[NAMModel] = []
        for f in sorted(self._models_dir.glob("*.nam")):
            try:
                with open(f) as fp:
                    data = json.load(fp)
                size_mb = f.stat().st_size / (1024 * 1024)
                models.append(
                    NAMModel(
                        name=f.stem,
                        path=str(f),
                        size_mb=size_mb,
                        architecture=data.get("architecture", "?"),
                        receptive_field=data.get("config", {}).get("receptive_field", 0),
                    )
                )
            except (json.JSONDecodeError, OSError) as exc:
                logger.debug("Skipping %s: %s", f.name, exc)
        return models

    # ── Crossfade helper (used by pipeline) ─────────────────────────

    def apply_crossfade(self, buf: np.ndarray) -> np.ndarray:
        """Apply a smooth crossfade from the model-switch boundary.

        Intended to be called by the pipeline after ``process()`` when
        a crossfade is active.
        """
        if self._crossfade_buf is None or self._crossfade_pos >= self._crossfade_len:
            return buf

        remain = min(self._crossfade_len - self._crossfade_pos, len(buf))
        fade_out = np.cos(np.linspace(0, np.pi / 2, remain)) ** 2
        fade_in = np.sin(np.linspace(0, np.pi / 2, remain)) ** 2

        buf[:remain] = (
            fade_out * self._crossfade_buf[:remain]
            + fade_in * buf[:remain]
        )

        self._crossfade_pos += remain
        return buf


# ── Standalone helpers ────────────────────────────────────────────────────


def available_models(models_dir: str | Path = DEFAULT_NAM_DIR) -> list[dict]:
    """Lightweight model listing — returns dicts (lighter than NAMModel).

    Returns
    -------
    list[dict]
        Each entry: ``{name, path, architecture, size_mb, feather, receptive_field}``.
    """
    host = NAMHost(models_dir=models_dir)
    results: list[dict] = []
    for m in host.list_available_models():
        results.append({
            "name": m.name,
            "path": m.path,
            "architecture": m.architecture,
            "size_mb": round(m.size_mb, 2),
            "feather": m.size_mb < 10.0,
            "receptive_field": m.receptive_field,
        })
    return results


def process_with_model(
    model_path: str,
    audio: np.ndarray,
    device: str = "cpu",
) -> np.ndarray:
    """Convenience: load a model, process one block, return audio."""
    host = NAMHost(device=device)
    host.load_model(model_path)
    return host.process(audio)
