# =============================================================================
#  Project   : modelmap
#  File      : src/modelmap/probes/linear.py
#  Purpose   : Linear probes with the mandatory controls doctrine (expA core)
#  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)
#  License   : All rights reserved (research code)
# =============================================================================
"""Linear probing with controls baked in, never optional (notes §4.1/§4.8).

Every probe run returns task accuracy AND its shuffled-label control on the
same split, so selectivity is always computable. Multinomial logistic
regression via plain NumPy gradient descent — no sklearn dependency, fully
deterministic given a seed, runs on any Mac.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass
class ProbeResult:
    task_accuracy: float
    control_accuracy: float      # shuffled-label control, same capacity/split
    selectivity: float
    n_train: int
    n_test: int
    seed: int


def _train_logreg(
    x: np.ndarray, y: np.ndarray, n_classes: int,
    l2: float, lr: float, epochs: int, rng: np.random.Generator,
) -> np.ndarray:
    """Multinomial logistic regression, full-batch GD. Returns (d+1, C)."""
    n, d = x.shape
    xb = np.hstack([x, np.ones((n, 1), dtype=x.dtype)])
    w = rng.normal(0, 0.01, size=(d + 1, n_classes)).astype(np.float64)
    onehot = np.eye(n_classes)[y]
    for _ in range(epochs):
        logits = xb @ w
        logits -= logits.max(axis=1, keepdims=True)
        p = np.exp(logits)
        p /= p.sum(axis=1, keepdims=True)
        grad = xb.T @ (p - onehot) / n + l2 * w
        w -= lr * grad
    return w


def _accuracy(w: np.ndarray, x: np.ndarray, y: np.ndarray) -> float:
    xb = np.hstack([x, np.ones((x.shape[0], 1), dtype=x.dtype)])
    return float((np.argmax(xb @ w, axis=1) == y).mean())


def probe_with_control(
    activations: np.ndarray,
    labels: np.ndarray,
    seed: int,
    train_frac: float = 0.8,
    l2: float = 1e-3,
    lr: float = 0.5,
    epochs: int = 300,
) -> ProbeResult:
    """Train task probe + shuffled-label control with identical capacity/split.

    The control shuffles TRAINING labels only (a probe that can still score
    high is memorizing, not reading structure); control test accuracy is
    evaluated against the true test labels, per Hewitt & Liang.
    """
    x = np.asarray(activations, dtype=np.float64)
    y = np.asarray(labels)
    classes, y_idx = np.unique(y, return_inverse=True)
    n_classes = len(classes)
    if n_classes < 2:
        raise ValueError("probe needs >=2 classes")

    rng = np.random.default_rng(seed)
    # standardize features on train split only
    perm = rng.permutation(len(y_idx))
    n_train = int(train_frac * len(y_idx))
    tr, te = perm[:n_train], perm[n_train:]
    mu, sd = x[tr].mean(0), x[tr].std(0) + 1e-8
    xs = (x - mu) / sd

    w_task = _train_logreg(xs[tr], y_idx[tr], n_classes, l2, lr, epochs, rng)
    task_acc = _accuracy(w_task, xs[te], y_idx[te])

    # NB: y_shuffled[tr] would be a fancy-indexed COPY — shuffle explicitly.
    shuffled_train = y_idx[tr].copy()
    rng.shuffle(shuffled_train)
    w_ctrl = _train_logreg(xs[tr], shuffled_train, n_classes, l2, lr, epochs, rng)
    ctrl_acc = _accuracy(w_ctrl, xs[te], y_idx[te])

    return ProbeResult(
        task_accuracy=task_acc,
        control_accuracy=ctrl_acc,
        selectivity=task_acc - ctrl_acc,
        n_train=len(tr),
        n_test=len(te),
        seed=seed,
    )