# =============================================================================
#  Project   : modelmap
#  File      : src/modelmap/capture/mlx_capture.py
#  Purpose   : Standard MLX capture layer — Tap wrappers over decoder layers
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Website   : https://modelmap.io
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""MLX activation capture (promoted from expH run #3, where retention
measured 1.004x plain prefill on a real 4-bit checkpoint).

Works on quantized mlx-lm models — the §4.9 capability gap. Layers are
wrapped with Tap objects (mlx-lm decoder loops call layers positionally, so
a plain callable suffices); d_model is inferred from runtime activations
because quantized embeddings pack their weight shapes; bf16 activations are
cast in MLX before NumPy conversion (numpy has no bfloat16).
"""

from __future__ import annotations

import numpy as np


class Tap:
    """Wraps a decoder layer; optionally retains its output or skips the block.

    skip=True implements layer ablation (the block's contribution is removed;
    the residual stream passes through unchanged) — the Gromov/ShortGPT-style
    intervention used by expC for causal verification.
    """

    def __init__(self, layer):
        self.layer = layer
        self.retained = None
        self.enabled = False
        self.skip = False
        self.edit = None      # optional callable applied to the block output

    def __call__(self, x, *args, **kwargs):
        if self.skip:
            return x
        out = self.layer(x, *args, **kwargs)
        if self.edit is not None:
            out = self.edit(out)
        if self.enabled:
            self.retained = out
        return out

    def __getattr__(self, name):  # delegate (e.g. .self_attn) for introspection
        return getattr(self.layer, name)


def install_taps(model) -> list[Tap]:
    """Wrap every decoder layer of an mlx-lm model in place; return the taps."""
    taps = [Tap(layer) for layer in model.model.layers]
    model.model.layers = taps
    return taps


def capture_mean_pooled(model, taps, token_ids_list, dtype=np.float32) -> np.ndarray:
    """Per-prompt, per-layer mean-pooled residual representations.

    Returns array of shape (n_prompts, n_layers, d_model). Batch = 1 per
    prompt (variable lengths, no padding confound); mean over sequence
    positions. Each forward is a full prefill (no KV cache reuse across
    prompts).
    """
    import mlx.core as mx

    for t in taps:
        t.enabled = True
    reps = None
    for i, ids in enumerate(token_ids_list):
        x = mx.array([list(ids)])
        model(x)
        pooled = [t.retained.mean(axis=1)[0].astype(mx.float32) for t in taps]
        mx.eval(*pooled)
        if reps is None:
            reps = np.zeros((len(token_ids_list), len(taps), pooled[0].shape[-1]), dtype=dtype)
        for li, p in enumerate(pooled):
            reps[i, li] = np.array(p, copy=False)
        for t in taps:
            t.retained = None
    for t in taps:
        t.enabled = False
    return reps


def capture_pooled(model, taps, token_ids_list, dtype=np.float32) -> dict[str, np.ndarray]:
    """Mean-pooled AND last-token reps in one pass.

    Returns {"mean": (n, L, d), "last": (n, L, d)} — added for expA run #2,
    where pooling choice is itself a measured variable.
    """
    import mlx.core as mx

    for t in taps:
        t.enabled = True
    mean_r = last_r = None
    for i, ids in enumerate(token_ids_list):
        x = mx.array([list(ids)])
        model(x)
        pooled = [t.retained.mean(axis=1)[0].astype(mx.float32) for t in taps]
        lasts = [t.retained[0, -1, :].astype(mx.float32) for t in taps]
        mx.eval(*pooled, *lasts)
        if mean_r is None:
            d = pooled[0].shape[-1]
            mean_r = np.zeros((len(token_ids_list), len(taps), d), dtype=dtype)
            last_r = np.zeros_like(mean_r)
        for li in range(len(taps)):
            mean_r[i, li] = np.array(pooled[li], copy=False)
            last_r[i, li] = np.array(lasts[li], copy=False)
        for t in taps:
            t.retained = None
    for t in taps:
        t.enabled = False
    return {"mean": mean_r, "last": last_r}


def random_init_twin(model_path):
    """Instantiate the same architecture with random weights (fp16, no
    quantization) — the architecture-only null mandated by charter §8.4."""
    import mlx.core as mx
    from mlx_lm.utils import _get_classes, load_config

    cfg = load_config(model_path)
    cfg_clean = {k: v for k, v in cfg.items() if k != "quantization"}
    model_class, args_class = _get_classes(cfg)
    twin = model_class(args_class.from_dict(cfg_clean))
    twin.set_dtype(mx.float16)
    return twin