# =============================================================================
#  Project   : localvm-research
#  File      : experiments/candidate_01/implementation/streaming_verifier.py
#  Purpose   : Layer-streamed q8 verification for models larger than free RAM —
#              per-layer materialize → compute → re-lazify on unified memory
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""StreamingVerifier — runs a verification forward pass through a quantized
model whose weights do NOT fit in free memory alongside the resident base.

Mechanism: the model is built with lazy (mmap-backed) weights. During a
forward pass we walk the layers manually; each layer's weights materialize
on first use, and immediately after the layer's output is evaluated we
re-assign that layer's parameters to FRESH lazy arrays (a new mx.load view),
dropping the concrete buffers. Peak residency ≈ resident base + a few
layers, while the SSD sees one sequential pass over the checkpoint per
sweep — exactly the expH-friendly access pattern.
"""

from __future__ import annotations

import glob
import time
from pathlib import Path

import mlx.core as mx
from mlx_lm import load as mlx_load
from mlx_lm.models.base import create_attention_mask


class StreamingVerifier:
    def __init__(self, model_path: str):
        self.path = Path(model_path)
        # lazy=True: parameters are mmap-backed lazy arrays, nothing evaluated
        self.model, self.tokenizer = mlx_load(str(model_path), lazy=True)
        self.shards = sorted(glob.glob(str(self.path / "*.safetensors")))
        self.weight_bytes = sum(Path(s).stat().st_size for s in self.shards)
        self.last_sweep_io_s = 0.0

    def _fresh_lazy_weights(self) -> dict:
        w = {}
        for s in self.shards:
            w.update(mx.load(s))  # lazy by default: no eval performed
        return w

    def _relazify(self, weights: dict, prefix: str) -> None:
        subset = [(k, v) for k, v in weights.items() if k.startswith(prefix)]
        if subset:
            self.model.load_weights(subset, strict=False)

    def forward_chunk(self, chunk_ids: list[int], cache) -> mx.array:
        """Teacher-force `chunk_ids` through the model with per-layer weight
        streaming. `cache` is a make_prompt_cache(self.model) list; it is
        advanced by len(chunk_ids). Returns logits (T, vocab)."""
        t0 = time.perf_counter()
        fresh = self._fresh_lazy_weights()
        inner = self.model.model
        x = mx.array(chunk_ids)[None]
        h = inner.embed_tokens(x)
        mx.eval(h)
        self._relazify(fresh, "model.embed_tokens")
        mask = create_attention_mask(h, cache)
        for i, layer in enumerate(inner.layers):
            h = layer(h, mask, cache=cache[i] if cache else None)
            mx.eval(h)
            self._relazify(fresh, f"model.layers.{i}.")
            if (i + 1) % 8 == 0:
                mx.clear_cache()  # release Metal allocator pools
        h = inner.norm(h)
        if hasattr(self.model, "lm_head"):
            logits = self.model.lm_head(h)
        else:  # tied embeddings
            logits = inner.embed_tokens.as_linear(h)
        logits = logits[0]
        mx.eval(logits)
        self._relazify(fresh, "model.norm")
        self._relazify(fresh, "lm_head")
        mx.clear_cache()
        self.last_sweep_io_s = time.perf_counter() - t0
        return logits