angular_error_shifts_sr.py

"""
angular_error_shifts_sr.py

Plot signed angular error distributions across r and s (counts, shared y-axis).
Uses Model(weight_var=7.0) and run_recall(rec_strength=r) with r in [0, 1.4].
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Any, List, Tuple
import copy

import numpy as np
import matplotlib

# Only force headless backend when this file is run as a script, not when imported
# (e.g. interactive_kappa_accessibility_explorer needs a GUI backend).
if __name__ == "__main__":
    matplotlib.use("Agg")
import matplotlib.pyplot as plt

from joblib import Parallel, delayed
from Model import Model
from PlaceInputs import PlaceInputs
from place_remapping import circular_shift_place_inputs

# ----------------------------
# Config
# ----------------------------
# Fast run: lighter bootstrap when QUICK_PREVIEW. Paper-scale: QUICK_PREVIEW False (1000 trials).
QUICK_PREVIEW = True

# Full 8×8 r/s grid (step 0.2) for 8-panel kappa-vs-accessibility plots under remapping.
# When True with PLACE_RECALL_SHIFT_RANDOM, outputs go to .../remap_rand_circular/fullgrid_n{N}/ .
RS_GRID_FULL = True

# Place p at encoding vs recall: use the same map (default), a fixed circular shift of
# recall p only, or an i.i.d. global shift drawn each trial (still sweep r,s over grid).
PLACE_RECALL_SHIFT_BINS = 0  # used when PLACE_RECALL_SHIFT_RANDOM is False
PLACE_RECALL_SHIFT_RANDOM = True  # False = same map encode/recall; True = random roll each trial

_remap_tag = (
    "remap_rand_circular"
    if PLACE_RECALL_SHIFT_RANDOM
    else (f"remap_shift_{PLACE_RECALL_SHIFT_BINS}" if PLACE_RECALL_SHIFT_BINS else "baseline")
)

_rs_full = np.round(np.arange(0.0, 1.41, 0.2), 2).tolist()
_rs_preview = [0.0, 0.4, 0.8, 1.2]
if RS_GRID_FULL:
    r_values = list(_rs_full)
    s_values = list(_rs_full)
elif QUICK_PREVIEW:
    r_values = list(_rs_preview)
    s_values = list(_rs_preview)
else:
    r_values = list(_rs_full)
    s_values = list(_rs_full)

if QUICK_PREVIEW:
    n_trials = 200
    BOOTSTRAP_N = 50
else:
    n_trials = 1000
    BOOTSTRAP_N = 200

if RS_GRID_FULL and QUICK_PREVIEW:
    _run_tag = f"fullgrid_n{n_trials}"
elif QUICK_PREVIEW:
    _run_tag = "preview"
else:
    _run_tag = "full"

OUTPUT_DIR = (
    Path(__file__).resolve().parent
    / "outputs"
    / "angular_error_shifts"
    / _remap_tag
    / _run_tag
)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

N_inp = 500
N_bar = 500
num_states = 100
steps = 30

weight_var_fixed = 7.0
seed_threshold = 0.5
num_caches_per_trial = 3
s_fixed_for_reg = 0.4
r_fixed_for_reg = 0.4
debug_single_trial = False
n_jobs = 2

# Inner breaks at -1.5, -0.5, 0.5, 1.5; outer edges at ±π so all signed angular errors are binned (none dropped).
ANG_ERR_HIST_BIN_EDGES = np.array([-np.pi, -1.5, -0.5, 0.5, 1.5, np.pi])


# ----------------------------
# Helpers
# ----------------------------
def signed_ring_error_states(pred: int, true: int, N: int) -> int:
    """Signed shortest-path error on ring in [-N/2, N/2)."""
    return int(((pred - true + N // 2) % N) - N // 2)


def decode_location(output_vec: np.ndarray, place_inputs: np.ndarray) -> int | None:
    if output_vec is None:
        return None
    if np.max(np.abs(output_vec)) < 1e-9:
        return None
    sims = place_inputs @ output_vec
    return int(np.argmax(sims))


def readout_sample_site(readout_norm: np.ndarray, rng: np.random.Generator) -> int:
    """Sample a retrieved site using normalized readout as probabilities."""
    if readout_norm is None or readout_norm.size == 0:
        return -1
    probs = np.clip(np.ravel(readout_norm.astype(float)), 0.0, None)
    denom = float(np.sum(probs))
    if denom <= 0 or not np.isfinite(denom):
        return -1
    probs = probs / denom
    return int(rng.choice(np.arange(probs.size), p=probs))


def nearest_cache(loc: int, cache_states: np.ndarray, N: int) -> int:
    dists = np.array([min(abs(loc - c), N - abs(loc - c)) for c in cache_states])
    return int(cache_states[int(np.argmin(dists))])


def retrieval_zone_width(
    decoded: np.ndarray,
    seed_out: np.ndarray,
    cache_states: np.ndarray,
    N: int,
    seed_thresh: float,
) -> float:
    """Mean #positions assigned to each cache (exclusive, width in sites)."""
    ok = seed_out > seed_thresh
    counts = {int(c): 0 for c in cache_states}
    for idx in range(N):
        if not ok[idx]:
            continue
        loc_hat = decoded[idx]
        if loc_hat < 0:
            continue
        # Assign each position to exactly one cache (nearest to decoded)
        assigned = nearest_cache(int(loc_hat), cache_states, N)
        counts[int(assigned)] += 1
    return float(np.mean(list(counts.values())))


def estimate_von_mises_kappa(angles: np.ndarray) -> float:
    """Estimate von Mises concentration kappa from samples in radians."""
    angles = np.asarray(angles)
    angles = angles[~np.isnan(angles)]
    if angles.size < 10:
        return float("nan")
    C = np.mean(np.cos(angles))
    S = np.mean(np.sin(angles))
    R = np.sqrt(C * C + S * S)
    if R < 1e-8:
        return 0.0
    if R < 0.53:
        kappa = 2 * R + R**3 + (5 * R**5) / 6
    elif R < 0.85:
        kappa = -0.4 + 1.39 * R + 0.43 / (1 - R)
    else:
        kappa = 1 / (R**3 - 4 * R**2 + 3 * R)
    n = angles.size
    if n < 15 and kappa > 0:
        kappa = max(kappa - 2 / (n * kappa), 0.0)
    return float(kappa)


def bootstrap_se(values: np.ndarray, stat_fn, n_boot: int = 200, seed: int = 0) -> float:
    values = np.asarray(values, dtype=float)
    values = values[np.isfinite(values)]
    if values.size < 2:
        return float("nan")
    rng = np.random.default_rng(seed)
    n = values.size
    boot = np.empty(n_boot, dtype=float)
    for i in range(n_boot):
        sample = values[rng.integers(0, n, size=n)]
        boot[i] = stat_fn(sample)
    return float(np.nanstd(boot, ddof=1))


def simulate_combo(
    r: float,
    s: float,
    place_inputs: np.ndarray,
    base_model: Model,
    *,
    recall_place_shift_bins: int = 0,
    recall_place_shift_random: bool = False,
    remap_prob: float | None = None,
    n_trials_override: int | None = None,
    j_sx_scale: float = 1.0,
    r_encode: float | None = None,
    r_recall: float | None = None,
    place_gain: float = 1.0,
) -> Dict[str, Any]:
    """
    place_inputs is the encoding map p_enc (caching + ΔJ_yx). If recall shifts are set,
    recall dynamics use p_recall = roll(p_enc, shift); decoding stays on p_enc so readout
    is interpreted in the learned coordinate frame.

    ``j_sx_scale`` scales the seed direction ``J_sx`` (Fang ``𝐣_{sx}``) in ``Model.run``; default 1.0.

    Fang notation: ``r`` → ``rec_strength``, ``s`` → ``search_factor``; place rows → ``place_inputs``
    (𝐩). If ``remap_prob`` is set (0..1), each trial uses a random coherent roll of recall **𝐩**
    with that probability, else recall **𝐩** matches encoding (ignores
    ``recall_place_shift_random`` / ``recall_place_shift_bins`` for the shift draw).

    If ``r_encode`` and/or ``r_recall`` are set, they override ``rec_strength`` for
    ``run_recurrent`` (encoding/caching) and ``run_recall`` respectively; omitted sides
    use ``r``.

    ``place_gain`` scales the place rows fed into dynamics (encoding, plasticity update,
    and recall), i.e.\ a gain on the feedforward place drive \(g\,\mathbf{p}\) with fixed
    ``J_{ix}``. Decoding still uses unshifted ``place_inputs`` as the coordinate frame.
    """
    re = float(r_encode) if r_encode is not None else float(r)
    rr = float(r_recall) if r_recall is not None else float(r)
    g = float(place_gain)
    rng = np.random.default_rng(
        int(1_000 * re + 10_000 * rr + 100 * s + 1_000_000 * int(round(g * 1_000))) + 321
    )
    errors_seeded: List[float] = []
    errors_nearest: List[float] = []
    width_sites_all: List[float] = []
    n_inp = int(place_inputs.shape[1])
    n_tr = int(n_trials_override) if n_trials_override is not None else n_trials

    for _ in range(n_tr):
        cache_states = rng.choice(np.arange(num_states), size=min(num_caches_per_trial, num_states), replace=False)

        if remap_prob is not None:
            rp = float(remap_prob)
            if rng.random() < rp:
                shift_bins = int(rng.integers(0, n_inp))
            else:
                shift_bins = 0
        elif recall_place_shift_random:
            shift_bins = int(rng.integers(0, n_inp))
        else:
            shift_bins = int(recall_place_shift_bins)
        place_recall = (
            circular_shift_place_inputs(place_inputs, shift_bins)
            if shift_bins != 0
            else place_inputs
        )
        place_enc = g * place_inputs
        place_rec = g * place_recall

        # Trial uses a fresh copy of the same base network (avoid network-to-network variability)
        model = copy.deepcopy(base_model)
        preacts, acts, _, _ = model.run_recurrent(place_enc, rec_strength=re)
        for cache_state in cache_states:
            model.update(place_enc[cache_state], acts[cache_state], preacts[cache_state])

        _, _, outputs, _ = model.run_recall(
            s, place_rec, rec_strength=rr, j_sx_scale=float(j_sx_scale)
        )
        place_output, seed_output = outputs
        seed_output = np.squeeze(seed_output)
        # Normalize seed output by max for thresholding
        seed_max = float(np.max(seed_output)) if np.max(seed_output) > 0 else 1.0
        seed_output_norm = seed_output / seed_max

        # Decode all positions (for accessibility width)
        decoded = np.full(num_states, -1, dtype=int)
        for bird_loc in range(num_states):
            loc_hat = decode_location(place_output[bird_loc], place_inputs)
            if loc_hat is None:
                continue
            decoded[bird_loc] = loc_hat

        width_sites = retrieval_zone_width(decoded, seed_output_norm, cache_states, num_states, seed_threshold)
        if debug_single_trial:
            print("DEBUG width_sites", width_sites, "num_states", num_states,
                  "cache_states", cache_states.tolist(), "seed_thresh", seed_threshold)
        width_sites_all.append(width_sites)

        # Random query location for error distribution
        bird_loc = int(rng.integers(0, num_states))
        if seed_output_norm[bird_loc] <= seed_threshold:
            continue

        # Sample retrieved site from normalized seed readout (single site)
        loc_hat = readout_sample_site(seed_output_norm, rng)
        if loc_hat < 0:
            continue

        # Define the target cache for this query as the nearest cache to bird_loc
        target_cache = nearest_cache(bird_loc, cache_states, num_states)

        # seeded (target) cache error
        err_states_seeded = signed_ring_error_states(loc_hat, target_cache, num_states)
        errors_seeded.append(2 * np.pi * err_states_seeded / num_states)

        # nearest cache error (precision uses this, matching rs_kappa_accessibility_regressions.py)
        nearest = nearest_cache(loc_hat, cache_states, num_states)
        err_states_nearest = signed_ring_error_states(loc_hat, nearest, num_states)
        errors_nearest.append(2 * np.pi * err_states_nearest / num_states)

    accessibility_sites = float(np.mean(width_sites_all)) if width_sites_all else float("nan")
    # Keep accessibility in states (not radians) for clarity
    kappa = estimate_von_mises_kappa(np.array(errors_nearest, dtype=float))
    return {
        "r": r,
        "r_encode": re,
        "r_recall": rr,
        "place_gain": g,
        "s": s,
        "errors_seeded": np.array(errors_seeded, dtype=float),
        "errors_nearest": np.array(errors_nearest, dtype=float),
        "accessibility_sites": accessibility_sites,
        "width_sites_all": np.array(width_sites_all, dtype=float),
        "kappa": kappa,
    }


def plot_hist_grid(results_by_rs: Dict[Tuple[float, float], Dict[str, Any]], key: str, title: str, filename: str) -> Path:
    bin_edges = ANG_ERR_HIST_BIN_EDGES
    max_density = 0.0
    for r in r_values:
        for s in s_values:
            dens, _ = np.histogram(results_by_rs[(r, s)][key], bins=bin_edges, density=True)
            max_density = max(max_density, float(dens.max()) if len(dens) > 0 else 0.0)

    fig, axes = plt.subplots(len(r_values), len(s_values), figsize=(4 * len(s_values), 3 * len(r_values)), sharey=True)
    if len(r_values) == 1 and len(s_values) == 1:
        axes = np.array([[axes]])
    elif len(r_values) == 1:
        axes = axes.reshape(1, -1)
    elif len(s_values) == 1:
        axes = axes.reshape(-1, 1)

    for i, r in enumerate(r_values):
        for j, s in enumerate(s_values):
            ax = axes[i, j]
            errs = results_by_rs[(r, s)][key]
            ax.hist(errs, bins=bin_edges, density=True, color="steelblue", alpha=0.8, edgecolor="black", linewidth=0.3)
            ax.axvline(0.0, color="red", linestyle="--", linewidth=1)
            ax.set_xlim([-np.pi, np.pi])
            ax.set_ylim([0, max_density * 1.05 if max_density > 0 else 1.0])
            ax.set_title(f"r={r:.2f}, s={s:.2f}")
            if i == len(r_values) - 1:
                ax.set_xlabel("Signed angular error (radians)")
            if j == 0:
                ax.set_ylabel("Probability density")

    plt.suptitle(title, fontsize=14, fontweight="bold", y=1.02)
    plt.tight_layout()
    outpath = OUTPUT_DIR / filename
    plt.savefig(outpath, dpi=300, bbox_inches="tight")
    plt.close(fig)
    return outpath


def plot_accessibility_heatmap(access_grid: np.ndarray, title: str, filename: str, cbar_label: str) -> Path:
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(access_grid, origin="lower", aspect="auto")
    ax.set_xticks(np.arange(len(s_values)))
    ax.set_xticklabels([f"{v:.1f}" for v in s_values])
    ax.set_yticks(np.arange(len(r_values)))
    ax.set_yticklabels([f"{v:.1f}" for v in r_values])
    ax.set_xlabel("search_factor s")
    ax.set_ylabel("rec_strength r")
    ax.set_title(title)
    cbar = fig.colorbar(im, ax=ax)
    cbar.set_label(cbar_label)
    outpath = OUTPUT_DIR / filename
    plt.savefig(outpath, dpi=300, bbox_inches="tight")
    plt.close(fig)
    return outpath


def linear_fit(x: np.ndarray, y: np.ndarray) -> Tuple[float, float, float, float]:
    """Return slope, intercept, r2, stderr_slope for simple linear fit."""
    x = np.asarray(x, dtype=float)
    y = np.asarray(y, dtype=float)
    n = x.size
    if n < 2 or not np.all(np.isfinite(y)):
        return float("nan"), float("nan"), float("nan"), float("nan")
    slope, intercept = np.polyfit(x, y, 1)
    y_hat = slope * x + intercept
    ss_res = np.sum((y - y_hat) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    # standard error of slope
    s_err = np.sqrt(ss_res / (n - 2)) if n > 2 else float("nan")
    sxx = np.sum((x - np.mean(x)) ** 2)
    stderr_slope = s_err / np.sqrt(sxx) if sxx > 0 else float("nan")
    return float(slope), float(intercept), float(r2), float(stderr_slope)


def interaction_fit(r_grid: np.ndarray, s_grid: np.ndarray, y_grid: np.ndarray) -> Tuple[np.ndarray, float]:
    """Fit y ~ b0 + b1*r + b2*s + b3*r*s over full grid; return coeffs and R2."""
    r_flat = r_grid.ravel()
    s_flat = s_grid.ravel()
    y_flat = y_grid.ravel()
    mask = np.isfinite(y_flat)
    r_flat = r_flat[mask]
    s_flat = s_flat[mask]
    y_flat = y_flat[mask]
    X = np.column_stack([np.ones_like(r_flat), r_flat, s_flat, r_flat * s_flat])
    coeffs, _, _, _ = np.linalg.lstsq(X, y_flat, rcond=None)
    y_hat = X @ coeffs
    ss_res = np.sum((y_flat - y_hat) ** 2)
    ss_tot = np.sum((y_flat - np.mean(y_flat)) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    return coeffs, float(r2)


def main() -> None:
    print("=" * 80)
    print("Angular error shifts (mean signed error)")
    print("=" * 80)
    print(f"QUICK_PREVIEW={QUICK_PREVIEW}")
    print(f"PLACE_RECALL_SHIFT_RANDOM={PLACE_RECALL_SHIFT_RANDOM}  PLACE_RECALL_SHIFT_BINS={PLACE_RECALL_SHIFT_BINS}")
    print(f"r values: {r_values}")
    print(f"s values: {s_values}")
    print(f"n_trials: {n_trials}  BOOTSTRAP_N={BOOTSTRAP_N}")
    print(f"OUTPUT_DIR: {OUTPUT_DIR}")

    place_inputs = PlaceInputs(N_inp, num_states).get_inputs()
    # One base network per sigma (weight_var), reused across all (r, s)
    base_model = Model(N_inp, N_bar, num_states, steps=steps, weight_var=weight_var_fixed)

    results_by_rs: Dict[Tuple[float, float], Dict[str, Any]] = {}
    accessibility_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)
    kappa_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)
    accessibility_se_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)
    kappa_se_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)
    variance_seeded_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)
    variance_nearest_grid = np.full((len(r_values), len(s_values)), np.nan, dtype=float)

    param_list = [(float(r), float(s)) for r in r_values for s in s_values]
    results = Parallel(n_jobs=n_jobs, verbose=5)(
        delayed(simulate_combo)(
            r,
            s,
            place_inputs,
            base_model,
            recall_place_shift_bins=PLACE_RECALL_SHIFT_BINS,
            recall_place_shift_random=PLACE_RECALL_SHIFT_RANDOM,
        )
        for (r, s) in param_list
    )
    for res in results:
        r = res["r"]
        s = res["s"]
        i = r_values.index(r)
        j = s_values.index(s)
        results_by_rs[(r, s)] = res
        accessibility_grid[i, j] = res["accessibility_sites"]
        kappa_grid[i, j] = res["kappa"]
        accessibility_se_grid[i, j] = bootstrap_se(
            res["width_sites_all"],
            np.mean,
            n_boot=BOOTSTRAP_N,
            seed=int(10_000 * r + 100 * s + 7),
        )
        kappa_se_grid[i, j] = bootstrap_se(
            res["errors_nearest"],
            estimate_von_mises_kappa,
            n_boot=BOOTSTRAP_N,
            seed=int(10_000 * r + 100 * s + 19),
        )
        errs_seeded = np.asarray(res["errors_seeded"], dtype=float)
        errs_nearest = np.asarray(res["errors_nearest"], dtype=float)
        variance_seeded_grid[i, j] = np.nanvar(errs_seeded) if errs_seeded.size else float("nan")
        variance_nearest_grid[i, j] = np.nanvar(errs_nearest) if errs_nearest.size else float("nan")

    hist_seeded = plot_hist_grid(
        results_by_rs,
        key="errors_seeded",
        title="Signed Angular Error (given retrieval; target cache; counts)",
        filename="signed_error_hist_grid_counts_seeded.png",
    )
    hist_nearest = plot_hist_grid(
        results_by_rs,
        key="errors_nearest",
        title="Signed Angular Error (given retrieval; nearest cache; counts)",
        filename="signed_error_hist_grid_counts_nearest.png",
    )
    access_heat = plot_accessibility_heatmap(
        accessibility_grid,
        title="Accessibility (retrieval-zone width in states)",
        filename="accessibility_width_heatmap.png",
        cbar_label="Width (states)",
    )
    kappa_heat = plot_accessibility_heatmap(
        kappa_grid,
        title="Precision (von Mises kappa)",
        filename="kappa_heatmap.png",
        cbar_label="Kappa",
    )

    x_r = np.array(r_values, dtype=float)
    x_s = np.array(s_values, dtype=float)
    r_grid, s_grid = np.meshgrid(x_r, x_s, indexing="ij")
    k_flat = kappa_grid.ravel()
    a_flat = accessibility_grid.ravel()
    r_flat = r_grid.ravel()
    _mask_panel = np.isfinite(k_flat) & np.isfinite(a_flat)
    k_flat = k_flat[_mask_panel]
    a_flat = a_flat[_mask_panel]
    r_flat = r_flat[_mask_panel]

    # 4-panel summary (Ip, Ik, It vs r; Ik vs Ip colored by r)
    ip_mean_by_r = np.nanmean(accessibility_grid, axis=1)
    ik_mean_by_r = np.nanmean(kappa_grid, axis=1)
    ip_se_by_r = np.nanstd(accessibility_grid, axis=1, ddof=1) / np.sqrt(np.sum(np.isfinite(accessibility_grid), axis=1))
    ik_se_by_r = np.nanstd(kappa_grid, axis=1, ddof=1) / np.sqrt(np.sum(np.isfinite(kappa_grid), axis=1))

    # Composite It as mean of min-max normalized Ip and Ik (per grid)
    ip_min, ip_max = np.nanmin(accessibility_grid), np.nanmax(accessibility_grid)
    ik_min, ik_max = np.nanmin(kappa_grid), np.nanmax(kappa_grid)
    ip_norm = (accessibility_grid - ip_min) / (ip_max - ip_min) if ip_max > ip_min else accessibility_grid * 0.0
    ik_norm = (kappa_grid - ik_min) / (ik_max - ik_min) if ik_max > ik_min else kappa_grid * 0.0
    it_grid = 0.5 * (ip_norm + ik_norm)
    it_mean_by_r = np.nanmean(it_grid, axis=1)
    it_se_by_r = np.nanstd(it_grid, axis=1, ddof=1) / np.sqrt(np.sum(np.isfinite(it_grid), axis=1))

    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    axes[0].errorbar(x_r, ip_mean_by_r, yerr=ip_se_by_r, fmt="o-", capsize=3, color="tab:purple")
    axes[0].set_title("I_p (accessibility)")
    axes[0].set_xlabel("rec_strength r")
    axes[0].set_ylabel("Width (states)")
    axes[0].grid(True, alpha=0.3)

    axes[1].errorbar(x_r, ik_mean_by_r, yerr=ik_se_by_r, fmt="o-", capsize=3, color="tab:blue")
    axes[1].set_title("I_k (precision)")
    axes[1].set_xlabel("rec_strength r")
    axes[1].set_ylabel("von Mises kappa")
    axes[1].grid(True, alpha=0.3)

    axes[2].errorbar(x_r, it_mean_by_r, yerr=it_se_by_r, fmt="o-", capsize=3, color="tab:green")
    axes[2].set_title("I_t (composite)")
    axes[2].set_xlabel("rec_strength r")
    axes[2].set_ylabel("Normalized units")
    axes[2].grid(True, alpha=0.3)

    sc2 = axes[3].scatter(k_flat, a_flat, c=r_flat, cmap="viridis", vmin=float(np.min(r_values)), vmax=float(np.max(r_values)))
    axes[3].set_title("I_k vs I_p")
    axes[3].set_xlabel("von Mises kappa")
    axes[3].set_ylabel("Accessibility (width, states)")
    axes[3].grid(True, alpha=0.3)
    cbar2 = fig.colorbar(sc2, ax=axes[3])
    cbar2.set_label("rec_strength r")

    panel_path = OUTPUT_DIR / "summary_ip_ik_it.png"
    plt.tight_layout()
    plt.savefig(panel_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    # kappa vs accessibility across full grid (with error bars; same mask as panel scatter)
    k_flat_reg = kappa_grid.ravel()
    a_flat_reg = accessibility_grid.ravel()
    k_se_flat = kappa_se_grid.ravel()
    a_se_flat = accessibility_se_grid.ravel()
    r_flat_reg = r_grid.ravel()
    s_flat_reg = s_grid.ravel()
    mask = np.isfinite(k_flat_reg) & np.isfinite(a_flat_reg)
    k_flat = k_flat_reg[mask]
    a_flat = a_flat_reg[mask]
    k_se_flat = k_se_flat[mask]
    a_se_flat = a_se_flat[mask]
    r_flat = r_flat_reg[mask]
    s_flat = s_flat_reg[mask]
    slope, intercept, r2, se = linear_fit(k_flat, a_flat)
    fig, ax = plt.subplots(figsize=(5, 4))
    ax.errorbar(
        k_flat,
        a_flat,
        xerr=k_se_flat,
        yerr=a_se_flat,
        fmt="none",
        ecolor="0.6",
        elinewidth=1,
        capsize=2,
        zorder=1,
    )
    sc = ax.scatter(
        k_flat,
        a_flat,
        c=r_flat,
        cmap="viridis",
        vmin=float(np.min(r_values)),
        vmax=float(np.max(r_values)),
        label="data",
        zorder=2,
    )
    for k, a, sv in zip(k_flat, a_flat, s_flat):
        ax.annotate(
            f"s={sv:.2f}",
            (k, a),
            xytext=(5, 5),
            textcoords="offset points",
            fontsize=7,
            color="0.15",
            zorder=3,
        )
    ax.plot(k_flat, slope * k_flat + intercept, color="tab:red", label="best fit")
    ax.set_xlabel("von Mises kappa")
    ax.set_ylabel("Accessibility (width, states)")
    ax.set_title("kappa vs accessibility")
    ax.grid(True, alpha=0.3)
    ax.legend()
    cbar = fig.colorbar(sc, ax=ax)
    cbar.set_label("rec_strength r")
    ax.text(0.02, 0.98, f"slope={slope:.2f}±{se:.2f}\nR²={r2:.2f}",
            transform=ax.transAxes, va="top", ha="left")
    kappa_access_path = OUTPUT_DIR / "regression_kappa_vs_accessibility.png"
    plt.savefig(kappa_access_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    variance_seeded_heat = plot_accessibility_heatmap(
        variance_seeded_grid,
        title="Variance of signed angular error (target cache)",
        filename="variance_seeded_heatmap.png",
        cbar_label="Variance (radians^2)",
    )
    variance_nearest_heat = plot_accessibility_heatmap(
        variance_nearest_grid,
        title="Variance of signed angular error (nearest cache)",
        filename="variance_nearest_heatmap.png",
        cbar_label="Variance (radians^2)",
    )

    # Linear regressions (all slices + interaction)
    print("\nLinear regressions (all slices):")
    # r vs kappa and r vs accessibility for each s
    for j, s in enumerate(s_values):
        y_k = kappa_grid[:, j]
        y_a = accessibility_grid[:, j]
        slope, intercept, r2, se = linear_fit(x_r, y_k)
        print(f"  r vs kappa (s={s:.1f}): slope={slope:.3f}±{se:.3f}, r2={r2:.3f}")
        slope, intercept, r2, se = linear_fit(x_r, y_a)
        print(f"  r vs accessibility (s={s:.1f}): slope={slope:.3f}±{se:.3f}, r2={r2:.3f}")
    print(f"  r values used: {r_values}")

    # s vs accessibility for each r
    for i, r in enumerate(r_values):
        y_a = accessibility_grid[i, :]
        slope, intercept, r2, se = linear_fit(x_s, y_a)
        print(f"  s vs accessibility (r={r:.1f}): slope={slope:.3f}±{se:.3f}, r2={r2:.3f}")
    print(f"  s values used: {s_values}")

    # Interaction tests across full grid
    coeffs_k, r2_k = interaction_fit(r_grid, s_grid, kappa_grid)
    coeffs_a, r2_a = interaction_fit(r_grid, s_grid, accessibility_grid)
    print("\nInteraction models (y ~ b0 + b1*r + b2*s + b3*r*s):")
    print(f"  kappa coeffs={coeffs_k}, r2={r2_k:.3f}")
    print(f"  accessibility coeffs={coeffs_a}, r2={r2_a:.3f}")

    # Plot regression: r (x) vs kappa (y) for fixed s
    if s_fixed_for_reg in s_values:
        s_idx = s_values.index(s_fixed_for_reg)
        y_k = kappa_grid[:, s_idx]
        slope, intercept, r2, se = linear_fit(x_r, y_k)
        fig, ax = plt.subplots(figsize=(5, 4))
        ax.errorbar(
            x_r,
            y_k,
            yerr=kappa_se_grid[:, s_idx],
            fmt="o",
            color="tab:blue",
            capsize=3,
            label="data",
        )
        ax.plot(x_r, slope * x_r + intercept, color="tab:red", label="best fit")
        ax.set_xlabel("rec_strength r")
        ax.set_ylabel("von Mises kappa")
        ax.set_title(f"r vs kappa (s={s_fixed_for_reg:.1f})")
        ax.grid(True, alpha=0.3)
        ax.legend()
        ax.text(0.02, 0.98, f"slope={slope:.2f}±{se:.2f}\nR²={r2:.2f}",
                transform=ax.transAxes, va="top", ha="left")
        reg_path = OUTPUT_DIR / "regression_r_vs_kappa.png"
        plt.savefig(reg_path, dpi=300, bbox_inches="tight")
        plt.close(fig)
        print(f"  {reg_path}")

        # r vs accessibility at fixed s
        y_a = accessibility_grid[:, s_idx]
        slope, intercept, r2, se = linear_fit(x_r, y_a)
        fig, ax = plt.subplots(figsize=(5, 4))
        ax.errorbar(
            x_r,
            y_a,
            yerr=accessibility_se_grid[:, s_idx],
            fmt="o",
            color="tab:blue",
            capsize=3,
            label="data",
        )
        ax.plot(x_r, slope * x_r + intercept, color="tab:red", label="best fit")
        ax.set_xlabel("rec_strength r")
        ax.set_ylabel("Accessibility (width, states)")
        ax.set_title(f"r vs accessibility (s={s_fixed_for_reg:.1f})")
        ax.grid(True, alpha=0.3)
        ax.legend()
        ax.text(0.02, 0.98, f"slope={slope:.2f}±{se:.2f}\nR²={r2:.2f}",
                transform=ax.transAxes, va="top", ha="left")
        reg_path = OUTPUT_DIR / "regression_r_vs_accessibility.png"
        plt.savefig(reg_path, dpi=300, bbox_inches="tight")
        plt.close(fig)
        print(f"  {reg_path}")

    # s vs accessibility at fixed r
    if r_fixed_for_reg in r_values:
        r_idx = r_values.index(r_fixed_for_reg)
        x_s = np.array(s_values, dtype=float)
        y_a = accessibility_grid[r_idx, :]
        slope, intercept, r2, se = linear_fit(x_s, y_a)
        fig, ax = plt.subplots(figsize=(5, 4))
        ax.errorbar(
            x_s,
            y_a,
            yerr=accessibility_se_grid[r_idx, :],
            fmt="o",
            color="tab:blue",
            capsize=3,
            label="data",
        )
        ax.plot(x_s, slope * x_s + intercept, color="tab:red", label="best fit")
        ax.set_xlabel("search_factor s")
        ax.set_ylabel("Accessibility (width, states)")
        ax.set_title(f"s vs accessibility (r={r_fixed_for_reg:.1f})")
        ax.grid(True, alpha=0.3)
        ax.legend()
        ax.text(0.02, 0.98, f"slope={slope:.2f}±{se:.2f}\nR²={r2:.2f}",
                transform=ax.transAxes, va="top", ha="left")
        reg_path = OUTPUT_DIR / "regression_s_vs_accessibility.png"
        plt.savefig(reg_path, dpi=300, bbox_inches="tight")
        plt.close(fig)
        print(f"  {reg_path}")

    save_path = OUTPUT_DIR / "angular_error_shifts.npz"
    np.savez_compressed(
        save_path,
        r_values=np.array(r_values),
        s_values=np.array(s_values),
        n_trials=n_trials,
        quick_preview=QUICK_PREVIEW,
        rs_grid_full=RS_GRID_FULL,
        bootstrap_n=BOOTSTRAP_N,
        weight_var_fixed=weight_var_fixed,
        num_states=num_states,
        seed_threshold=seed_threshold,
        num_caches_per_trial=num_caches_per_trial,
        place_recall_shift_bins=PLACE_RECALL_SHIFT_BINS,
        place_recall_shift_random=PLACE_RECALL_SHIFT_RANDOM,
        results_by_rs=results_by_rs,
        accessibility_grid=accessibility_grid,
        kappa_grid=kappa_grid,
        accessibility_se_grid=accessibility_se_grid,
        kappa_se_grid=kappa_se_grid,
        variance_seeded_grid=variance_seeded_grid,
        variance_nearest_grid=variance_nearest_grid,
    )

    print("Saved:")
    print(f"  {hist_seeded}")
    print(f"  {hist_nearest}")
    print(f"  {access_heat}")
    print(f"  {kappa_heat}")
    print(f"  {kappa_access_path}")
    print(f"  {panel_path}")
    print(f"  {variance_seeded_heat}")
    print(f"  {variance_nearest_heat}")
    print(f"  {save_path}")


if __name__ == "__main__":
    main()