code/training/evaluate.py

"""
Evaluation module for the Money Experiment.

Loads trained Model A and Model B checkpoints, runs inference on
Allen-matched circuit configurations, and compares predictions to
real Allen Neuropixels statistics.

The key metric: Does Model B (neuromod-aware) predict the
running-vs-stationary DIFFERENCE better than Model A (plain HH)?
"""

from __future__ import annotations

import glob
import json
import logging
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

from .config import INPUT_FEATURES_A, INPUT_FEATURES_B, OUTPUT_STATS, TrainConfig
from .dataset import Normalizer
from .model import CircuitTransformer

logger = logging.getLogger(__name__)


# ── Load Allen data ──────────────────────────────────────────────────────────

def load_allen_epochs(allen_dir: str) -> dict[str, list[dict]]:
    """Load Allen epoch JSONs grouped by (session_id, state).

    Returns:
        {
            "running": [{"session_id": ..., "statistics": {...}}, ...],
            "stationary": [{"session_id": ..., "statistics": {...}}, ...],
        }
    """
    pattern = str(Path(allen_dir) / "allen_*.json")
    files = sorted(glob.glob(pattern))
    if not files:
        raise FileNotFoundError(f"No Allen epoch files found in {allen_dir}")

    running = []
    stationary = []
    for fpath in files:
        with open(fpath) as f:
            data = json.load(f)
        state = data.get("epoch_type", data.get("behavioral_state", "unknown"))
        if state == "running":
            running.append(data)
        elif state == "stationary":
            stationary.append(data)

    logger.info(
        f"Allen data: {len(running)} running epochs, "
        f"{len(stationary)} stationary epochs from {len(files)} files"
    )
    return {"running": running, "stationary": stationary}


def allen_sessions_with_both_states(
    allen_data: dict[str, list[dict]],
) -> list[int]:
    """Find session IDs that have both running AND stationary epochs."""
    running_sessions = {d["session_id"] for d in allen_data["running"]}
    stationary_sessions = {d["session_id"] for d in allen_data["stationary"]}
    both = sorted(running_sessions & stationary_sessions)
    logger.info(f"Sessions with both states: {len(both)}")
    return both


def compute_allen_session_means(
    allen_data: dict[str, list[dict]],
    session_ids: list[int],
) -> dict[str, dict[int, dict[str, float]]]:
    """Compute mean statistics per session per state.

    Returns:
        {
            "running":    {session_id: {stat: mean_value, ...}, ...},
            "stationary": {session_id: {stat: mean_value, ...}, ...},
        }
    """
    result = {"running": {}, "stationary": {}}

    for state in ["running", "stationary"]:
        for sid in session_ids:
            epochs = [
                d for d in allen_data[state] if d["session_id"] == sid
            ]
            if not epochs:
                continue

            means = {}
            for stat in OUTPUT_STATS:
                vals = [e["statistics"][stat] for e in epochs if stat in e.get("statistics", {})]
                if vals:
                    means[stat] = float(np.mean(vals))
            result[state][sid] = means

    return result


# ── Load model from checkpoint ───────────────────────────────────────────────

def load_model_from_checkpoint(
    ckpt_path: str, device: torch.device
) -> tuple[nn.Module, Normalizer, Normalizer, dict]:
    """Load a trained model + normalizers from a checkpoint.

    Supports both Transformer and MLP checkpoints.

    Returns:
        (model, x_norm, y_norm, meta)
    """
    from .model import CircuitMLP

    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = ckpt["config"]
    arch = cfg.get("arch", "transformer")

    if arch == "mlp":
        model = CircuitMLP(
            n_features=cfg["n_features"],
            n_outputs=cfg["n_outputs"],
            hidden_dims=cfg.get("hidden_dims", [64, 64]),
            dropout=cfg.get("dropout", 0.1),
        )
    else:
        model = CircuitTransformer(
            n_features=cfg["n_features"],
            n_outputs=cfg["n_outputs"],
            d_model=cfg["d_model"],
            n_heads=cfg["n_heads"],
            n_layers=cfg["n_layers"],
            d_ff=cfg["d_ff"],
            dropout=cfg["dropout"],
            has_ach=cfg.get("has_ach", cfg["n_features"] > 10),
        )
    model.load_state_dict(ckpt["model_state_dict"])
    model = model.to(device)
    model.eval()

    x_norm = Normalizer.from_state_dict(ckpt["x_norm"])
    y_norm = Normalizer.from_state_dict(ckpt["y_norm"])
    meta = ckpt["meta"]

    logger.info(
        f"Loaded Model {meta['model_variant']} ({arch}) from {ckpt_path} "
        f"(epoch {ckpt['epoch']}, val_loss={ckpt['val_loss']:.5f})"
    )
    return model, x_norm, y_norm, meta


# ── Prediction helpers ───────────────────────────────────────────────────────

def build_allen_matched_input(
    session_means: dict[str, float],
    ach_level: float,
    model_variant: str,
) -> dict[str, float]:
    """Build an input feature dict that approximates an Allen V1 circuit.

    Since we don't know the exact circuit structure of the real brain,
    we use canonical values from our simulation parameter ranges:
      - n_exc=160, n_inh=40 (200 total, 80/20 split)
      - conn_prob=0.06 (canonical cortical)
      - n_synapses ≈ 200*200*0.06 = 2400
      - mean_in_degree ≈ 200*0.06 = 12
      - gS, OU params at canonical values
      - ACh: 0.0 for stationary (low ACh), 1.0 for running (high ACh)
    """
    # Canonical V1-like circuit params (median of our sim distribution)
    inp = {
        "n_exc": 160.0,
        "n_inh": 40.0,
        "conn_prob": 0.06,
        "n_synapses": 2400.0,
        "mean_in_degree": 12.0,
        "gS_exc_effective": 5e-6,  # baseline (ACh=0)
        "ou_mu_effective": -0.001,
        "ou_sigma_effective": 0.001,
        "ou_tau": 5.0,
        "sim_duration_ms": 3000.0,
    }

    if model_variant == "B":
        inp["ach_level"] = ach_level
        # Adjust gS_exc_effective for ACh modulation (E2 curve)
        import math
        syn_scale = max(0.05, math.exp(-2.3 * ach_level))
        inp["gS_exc_effective"] = 5e-6 * syn_scale

    return inp


@torch.no_grad()
def predict_statistics(
    model: CircuitTransformer,
    x_norm: Normalizer,
    y_norm: Normalizer,
    input_dict: dict[str, float],
    input_features: list[str],
    device: torch.device,
) -> dict[str, float]:
    """Run model inference on a single input → predicted statistics.

    Returns dict of {stat_name: predicted_value} in ORIGINAL scale.
    """
    # Build input vector
    x = np.array([[input_dict[f] for f in input_features]], dtype=np.float64)

    # Normalize
    x_n = x_norm.transform(x)
    x_t = torch.tensor(x_n, dtype=torch.float32).to(device)

    # Predict (normalized space)
    y_n = model(x_t).cpu().numpy()

    # Inverse normalize
    y = y_norm.inverse(y_n)

    # Build output dict
    return {stat: float(y[0, i]) for i, stat in enumerate(OUTPUT_STATS)}


# ── The Money Experiment ─────────────────────────────────────────────────────

def run_money_experiment(cfg: TrainConfig, device: torch.device | None = None) -> dict:
    """The core experiment: compare Model A vs Model B on real Allen data.

    Steps:
      1. Load Allen data, find sessions with both states
      2. Load both trained models
      3. For each session: predict stats at low-ACh + high-ACh
      4. Compare predicted deltas to observed deltas
      5. Compute transfer metrics

    Returns:
        Dict with all results for paper figures.
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    logger.info("=" * 70)
    logger.info("THE MONEY EXPERIMENT: Sim-to-Real Transfer")
    logger.info("=" * 70)

    # 1. Load Allen data
    allen_data = load_allen_epochs(cfg.allen_dir)
    session_ids = allen_sessions_with_both_states(allen_data)
    if not session_ids:
        raise ValueError("No sessions with both running and stationary data!")
    session_means = compute_allen_session_means(allen_data, session_ids)

    # 2. Load models
    ckpt_a = str(Path(cfg.checkpoint_dir) / "model_a" / "best.pt")
    ckpt_b = str(Path(cfg.checkpoint_dir) / "model_b" / "best.pt")

    model_a, x_norm_a, y_norm_a, meta_a = load_model_from_checkpoint(ckpt_a, device)
    model_b, x_norm_b, y_norm_b, meta_b = load_model_from_checkpoint(ckpt_b, device)

    # 3. Predictions
    results_per_session = {}

    for sid in session_ids:
        real_running = session_means["running"].get(sid, {})
        real_stationary = session_means["stationary"].get(sid, {})
        if not real_running or not real_stationary:
            continue

        # Model A: predict at ACh=0 for both states (it has no ACh concept)
        inp_a = build_allen_matched_input(real_stationary, ach_level=0.0, model_variant="A")
        pred_a = predict_statistics(
            model_a, x_norm_a, y_norm_a, inp_a, INPUT_FEATURES_A, device
        )

        # Model B: predict at ACh=0 (stationary) and ACh=1.0 (running)
        inp_b_low = build_allen_matched_input(real_stationary, ach_level=0.0, model_variant="B")
        inp_b_high = build_allen_matched_input(real_running, ach_level=1.0, model_variant="B")

        pred_b_low = predict_statistics(
            model_b, x_norm_b, y_norm_b, inp_b_low, INPUT_FEATURES_B, device
        )
        pred_b_high = predict_statistics(
            model_b, x_norm_b, y_norm_b, inp_b_high, INPUT_FEATURES_B, device
        )

        results_per_session[sid] = {
            "real_running": real_running,
            "real_stationary": real_stationary,
            "pred_a": pred_a,  # Model A: same prediction for both states
            "pred_b_low": pred_b_low,  # Model B @ ACh=0
            "pred_b_high": pred_b_high,  # Model B @ ACh=1
        }

    # 4. Compute transfer metrics
    logger.info(f"\nAnalyzing {len(results_per_session)} sessions...")

    # For each statistic, compute:
    #   - Real delta: running - stationary
    #   - Model A delta: 0 (can't distinguish states)
    #   - Model B delta: pred_high - pred_low
    #   - Correlation of predicted vs real deltas across sessions

    stat_metrics = {}
    for stat in OUTPUT_STATS:
        real_deltas = []
        pred_b_deltas = []
        pred_a_vals = []
        real_running_vals = []
        real_stationary_vals = []
        pred_b_high_vals = []
        pred_b_low_vals = []

        for sid, res in results_per_session.items():
            if stat in res["real_running"] and stat in res["real_stationary"]:
                real_r = res["real_running"][stat]
                real_s = res["real_stationary"][stat]
                real_deltas.append(real_r - real_s)
                real_running_vals.append(real_r)
                real_stationary_vals.append(real_s)
                pred_a_vals.append(res["pred_a"].get(stat, 0))
                pred_b_low_vals.append(res["pred_b_low"].get(stat, 0))
                pred_b_high_vals.append(res["pred_b_high"].get(stat, 0))
                pred_b_deltas.append(
                    res["pred_b_high"].get(stat, 0) - res["pred_b_low"].get(stat, 0)
                )

        if len(real_deltas) < 3:
            stat_metrics[stat] = {"n_sessions": len(real_deltas), "skip": True}
            continue

        real_deltas = np.array(real_deltas)
        pred_b_deltas = np.array(pred_b_deltas)

        # Correlation between predicted and real deltas (Model B)
        if np.std(real_deltas) > 0 and np.std(pred_b_deltas) > 0:
            delta_corr_b = float(np.corrcoef(real_deltas, pred_b_deltas)[0, 1])
        else:
            delta_corr_b = 0.0

        # Model A: correlation between single prediction and real running/stationary
        real_all = np.array(real_running_vals + real_stationary_vals)
        pred_a_all = np.array(pred_a_vals + pred_a_vals)  # same prediction twice
        if np.std(real_all) > 0 and np.std(pred_a_all) > 0:
            corr_a = float(np.corrcoef(real_all, pred_a_all)[0, 1])
        else:
            corr_a = 0.0

        # Model B: correlation between predictions and real values
        pred_b_all = np.array(pred_b_high_vals + pred_b_low_vals)
        if np.std(real_all) > 0 and np.std(pred_b_all) > 0:
            corr_b = float(np.corrcoef(real_all, pred_b_all)[0, 1])
        else:
            corr_b = 0.0

        # Sign accuracy: does Model B predict the correct direction of change?
        sign_correct = float(np.mean(np.sign(real_deltas) == np.sign(pred_b_deltas)))

        # Mean real delta
        mean_real_delta = float(np.mean(real_deltas))
        mean_pred_b_delta = float(np.mean(pred_b_deltas))

        stat_metrics[stat] = {
            "n_sessions": len(real_deltas),
            "delta_corr_b": round(delta_corr_b, 4),
            "overall_corr_a": round(corr_a, 4),
            "overall_corr_b": round(corr_b, 4),
            "sign_accuracy_b": round(sign_correct, 4),
            "mean_real_delta": round(mean_real_delta, 4),
            "mean_pred_b_delta": round(mean_pred_b_delta, 4),
        }

    # 5. Summary
    logger.info(f"\n{'='*70}")
    logger.info("MONEY EXPERIMENT RESULTS")
    logger.info(f"{'='*70}")
    logger.info(f"  {'Statistic':25s}  {'Corr A':>8s}  {'Corr B':>8s}  {'Δ Corr B':>10s}  {'Sign%':>6s}")
    logger.info(f"  {'-'*25}  {'-'*8}  {'-'*8}  {'-'*10}  {'-'*6}")

    mean_corr_a = []
    mean_corr_b = []
    mean_delta_corr = []

    for stat in OUTPUT_STATS:
        m = stat_metrics[stat]
        if m.get("skip"):
            logger.info(f"  {stat:25s}  SKIPPED (n={m['n_sessions']})")
            continue
        logger.info(
            f"  {stat:25s}  {m['overall_corr_a']:8.4f}  {m['overall_corr_b']:8.4f}  "
            f"{m['delta_corr_b']:10.4f}  {m['sign_accuracy_b']:6.1%}"
        )
        mean_corr_a.append(m["overall_corr_a"])
        mean_corr_b.append(m["overall_corr_b"])
        mean_delta_corr.append(m["delta_corr_b"])

    if mean_corr_a:
        logger.info(f"  {'-'*25}  {'-'*8}  {'-'*8}  {'-'*10}  {'-'*6}")
        logger.info(
            f"  {'MEAN':25s}  {np.mean(mean_corr_a):8.4f}  {np.mean(mean_corr_b):8.4f}  "
            f"{np.mean(mean_delta_corr):10.4f}"
        )
    logger.info(f"{'='*70}")

    # Key question answer
    b_wins = sum(1 for s in OUTPUT_STATS
                 if not stat_metrics[s].get("skip")
                 and stat_metrics[s]["overall_corr_b"] > stat_metrics[s]["overall_corr_a"])
    total_compared = sum(1 for s in OUTPUT_STATS if not stat_metrics[s].get("skip"))

    if total_compared > 0:
        logger.info(
            f"\n  KEY RESULT: Model B (ACh) beats Model A (plain) on "
            f"{b_wins}/{total_compared} statistics "
            f"({b_wins/total_compared:.0%})"
        )

    # Save full results
    experiment_results = {
        "n_sessions": len(results_per_session),
        "session_ids": list(results_per_session.keys()),
        "stat_metrics": stat_metrics,
        "summary": {
            "mean_corr_a": round(float(np.mean(mean_corr_a)), 4) if mean_corr_a else None,
            "mean_corr_b": round(float(np.mean(mean_corr_b)), 4) if mean_corr_b else None,
            "mean_delta_corr_b": round(float(np.mean(mean_delta_corr)), 4) if mean_delta_corr else None,
            "b_wins": b_wins,
            "total_compared": total_compared,
        },
        "per_session": {
            str(k): v for k, v in results_per_session.items()
        },
    }

    # Save to disk
    log_dir = Path(cfg.log_dir)
    log_dir.mkdir(parents=True, exist_ok=True)
    with open(log_dir / "money_experiment_results.json", "w") as f:
        json.dump(experiment_results, f, indent=2)
    logger.info(f"Full results saved to {log_dir / 'money_experiment_results.json'}")

    return experiment_results