scripts/memory_ablation.py

"""4-arm memory ablation across E2/E3/E4 to prove memory CAUSALLY changes outcomes.

Arms:
    A — keyword_security_triager (no memory queries — just heuristic)
    B — keyword_security_triager w/ memory store disabled (env returns no hits;
        OPSGUARD_MEMORY_DISABLED=1)
    C — memory_aware policy (current: periodic query_history)
    D — memory_aware + ContributorProfile features actively consulted in policy
        (already wired into observation; this arm uses them)

Outputs (under --out, default `eval_outputs/memory_ablation/`):
    rollouts.jsonl         per-rollout records
    summary.md             arm × scenario reward (mean ± std) + attack_recall
    delta.png              bar plot of arm reward delta vs Arm B (no-memory baseline)
    pmi.json               PMI between memory hit in last 5 steps and a correct
                           attack-detection action (counterfactual proxy)

Usage:
    python scripts/memory_ablation.py
    python scripts/memory_ablation.py --scenarios E2_social_eng_buildup --seeds 0 1
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import time
from pathlib import Path

# Make repo importable when run as a script.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from eval.harness import rollout
from eval.policies import build_policy, make_keyword_security_triager
from models import ActionType, OpsguardAction, OpsguardObservation, SecurityVerdict
from server.opsguard_environment import OpsguardEnvironment


DEFAULT_SCENARIOS = [
    "E2_social_eng_buildup",
    "E3_compromised_maintainer",
    "E4_multi_vector",
]
DEFAULT_SEEDS = [0, 1, 2]

CATCH_ACTIONS = {"reject_pr", "close_spam", "flag_security"}
PMI_LOOKBACK = 5


# ------------------------------------------------------------------------------
# Arm D: memory-aware policy that ACTUALLY consults ContributorProfile features.
# ------------------------------------------------------------------------------
def make_memory_aware_profile_triager():
    """Memory-aware triager + per-author behavioural change-point detection.

    Periodically issues `query_history` like the base memory_aware policy, and
    additionally inspects `current_issue.author_profile` to flag PRs whose
    behavioural delta or anomaly score sharply diverges from the contributor's
    established baseline (Jia-Tan-style change-point).
    """
    base = make_keyword_security_triager()
    state = {"last_query_step": -10}

    def act(obs: OpsguardObservation) -> OpsguardAction:
        ci = obs.current_issue
        if ci is None:
            return OpsguardAction(action_type=ActionType.WAIT)

        # Periodic memory query, same cadence as memory_aware.
        if obs.step > 0 and obs.step - state["last_query_step"] >= 8:
            state["last_query_step"] = obs.step
            return OpsguardAction(
                action_type=ActionType.QUERY_HISTORY,
                query=ci.author_login,
                reasoning="check prior author activity",
            )

        # Profile-driven change-point detection BEFORE falling through to the
        # keyword triager. A long-trusted author (high trust_score, several prior
        # PRs) suddenly producing a PR with a large behavioural delta or anomaly
        # spike is the supply-chain signal we want to catch.
        prof = ci.author_profile or {}
        n_seen = prof.get("n_prs_seen", 0)
        trust = prof.get("trust_score", 0.5)
        anomaly = prof.get("anomaly_score_ewma", 0.0)
        delta = prof.get("behavioral_delta", 0.0)

        if ci.is_pr and n_seen >= 3 and trust >= 0.7 and (anomaly >= 1.5 or delta >= 200.0):
            return OpsguardAction(
                action_type=ActionType.FLAG_SECURITY,
                target_issue_id=ci.issue_id,
                comment_body=(
                    f"behavioural change-point on trusted author "
                    f"(anomaly_ewma={anomaly:.2f}, delta={delta:.1f})"
                ),
                reasoning="contributor profile change-point",
                security_verdict=SecurityVerdict.SUSPICIOUS,
            )

        return base(obs)

    return act


# ------------------------------------------------------------------------------
# Arm specs.
# ------------------------------------------------------------------------------
ARMS = {
    "A_no_memory_query": {
        "policy_factory": lambda: build_policy("keyword_security_triager"),
        "env_overrides": {},
        "description": "keyword triager; never issues query_history",
    },
    "B_memory_disabled": {
        "policy_factory": lambda: build_policy("keyword_security_triager"),
        "env_overrides": {"OPSGUARD_MEMORY_DISABLED": "1"},
        "description": "keyword triager; env honors OPSGUARD_MEMORY_DISABLED → memory_hits always []",
    },
    "C_memory_aware": {
        "policy_factory": lambda: build_policy("memory_aware"),
        "env_overrides": {},
        "description": "memory_aware: periodic query_history",
    },
    "D_memory_aware_profile": {
        "policy_factory": make_memory_aware_profile_triager,
        "env_overrides": {},
        "description": "memory_aware + ContributorProfile change-point features",
    },
}


# ------------------------------------------------------------------------------
# Custom rollout that mirrors eval.harness.rollout but ALSO records the per-step
# trace we need for PMI: action, was_attack, was_caught, memory_hits_in_obs.
# ------------------------------------------------------------------------------
def rollout_with_trace(
    env: OpsguardEnvironment,
    policy_name: str,
    policy_fn,
    scenario_id: str,
    seed: int,
):
    t0 = time.time()
    obs = env.reset(scenario_id=scenario_id, seed=seed)
    cum = 0.0
    last_meta: dict = {}
    n_steps = 0
    trace: list[dict] = []
    budget = env._episode.scenario.step_budget + 5  # type: ignore[union-attr]

    while not obs.done and n_steps < budget:
        # Snapshot memory-hit count BEFORE acting (the obs the policy sees).
        mem_hits_now = len(obs.memory_hits or [])
        action = policy_fn(obs)
        target_id = action.target_issue_id
        if target_id is None and obs.current_issue is not None:
            target_id = obs.current_issue.issue_id

        obs = env.step(action)
        if obs.reward is not None:
            cum += obs.reward
        n_steps += 1
        if obs.metadata:
            last_meta = obs.metadata

        # The most recent action lands at the end of recent_actions; pull
        # ground-truth `is_attack` from there (env populated it).
        recent = (obs.recent_actions or [])
        last_rec = recent[-1] if recent else {}
        action_str = last_rec.get("action", action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type))
        was_attack = bool(last_rec.get("is_attack", False))
        was_catch_action = action_str in CATCH_ACTIONS
        # "Caught" = a catch action applied to a true attack.
        caught = was_catch_action and was_attack

        trace.append({
            "step": n_steps,
            "action": action_str,
            "issue_id": last_rec.get("issue_id", target_id),
            "is_attack": was_attack,
            "catch_action": was_catch_action,
            "caught": caught,
            "memory_hits_in_obs": mem_hits_now,
        })

    record = {
        "policy": policy_name,
        "scenario_id": scenario_id,
        "seed": seed,
        "cumulative_reward": round(cum, 4),
        "n_steps": n_steps,
        "n_resolved": last_meta.get("legit_resolved", 0),
        "n_total": last_meta.get("legit_total", 0),
        "n_spam_caught": last_meta.get("attacks_caught", 0),
        "n_spam_total": last_meta.get("attacks_total", 0),
        "elapsed_sec": round(time.time() - t0, 2),
        "final_breakdown": last_meta,
        "trace": trace,
    }
    return record


# ------------------------------------------------------------------------------
# Aggregation, PMI, summary, plot.
# ------------------------------------------------------------------------------
def _mean_std(xs: list[float]) -> tuple[float, float]:
    if not xs:
        return 0.0, 0.0
    m = sum(xs) / len(xs)
    if len(xs) < 2:
        return m, 0.0
    v = sum((x - m) ** 2 for x in xs) / (len(xs) - 1)
    return m, math.sqrt(v)


def aggregate_records(records: list[dict]) -> dict:
    cells = {}
    for r in records:
        cells.setdefault((r["policy"], r["scenario_id"]), []).append(r)
    out = []
    for (arm, scen), rs in sorted(cells.items()):
        rewards = [r["cumulative_reward"] for r in rs]
        recall = [
            r["n_spam_caught"] / r["n_spam_total"]
            for r in rs if r["n_spam_total"]
        ]
        rm, rs_ = _mean_std(rewards)
        recm, _recs = _mean_std(recall)
        out.append({
            "arm": arm,
            "scenario": scen,
            "n": len(rs),
            "reward_mean": round(rm, 3),
            "reward_std": round(rs_, 3),
            "attack_recall_mean": round(recm, 3) if recall else None,
            "spam_caught_mean": round(sum(r["n_spam_caught"] for r in rs) / len(rs), 2),
            "spam_total_mean": round(sum(r["n_spam_total"] for r in rs) / len(rs), 2),
        })
    return {"cells": out}


def compute_pmi(records: list[dict]) -> dict:
    """Per-arm PMI between "memory hit retrieved in prior 5 steps" and "this step
    is a correct attack-detection (catch action on a real attack)".

    For each catch-action step (action ∈ {reject_pr, close_spam, flag_security}),
    we ask:
        - Was there ANY memory hit observed in the prior PMI_LOOKBACK steps?
        - Did this catch a true attack?

    PMI_proxy = P(catch | retrieved_recent) / P(catch)
    Reported per arm. Values > 1 ⇒ memory retrieval covaries with successful
    attack detection above chance.
    """
    by_arm: dict[str, list[dict]] = {}
    for r in records:
        by_arm.setdefault(r["policy"], []).append(r)

    out: dict[str, dict] = {}
    for arm, recs in by_arm.items():
        # Pool catch-action events across all rollouts for this arm.
        catch_events: list[dict] = []  # one entry per catch action
        for r in recs:
            trace = r.get("trace", [])
            for i, step in enumerate(trace):
                if not step["catch_action"]:
                    continue
                window = trace[max(0, i - PMI_LOOKBACK):i]
                retrieved_recent = any(s["memory_hits_in_obs"] > 0 for s in window)
                catch_events.append({
                    "caught": bool(step["caught"]),
                    "retrieved_recent": retrieved_recent,
                })

        n = len(catch_events)
        n_caught = sum(1 for e in catch_events if e["caught"])
        n_retrieved = sum(1 for e in catch_events if e["retrieved_recent"])
        n_caught_and_retrieved = sum(
            1 for e in catch_events if e["caught"] and e["retrieved_recent"]
        )

        p_catch = n_caught / n if n else 0.0
        p_catch_given_retrieved = (
            n_caught_and_retrieved / n_retrieved if n_retrieved else 0.0
        )
        # PMI proxy as a ratio (paper-style "lift").
        if p_catch > 0 and n_retrieved > 0:
            pmi_ratio = p_catch_given_retrieved / p_catch
            pmi_log2 = math.log2(pmi_ratio) if pmi_ratio > 0 else float("-inf")
        else:
            pmi_ratio = None
            pmi_log2 = None

        out[arm] = {
            "n_catch_events": n,
            "n_caught": n_caught,
            "n_with_recent_memory_hit": n_retrieved,
            "n_caught_and_retrieved": n_caught_and_retrieved,
            "p_catch": round(p_catch, 4),
            "p_catch_given_retrieved": round(p_catch_given_retrieved, 4),
            "pmi_ratio_p_catch_given_retrieved_over_p_catch": (
                round(pmi_ratio, 4) if pmi_ratio is not None else None
            ),
            "pmi_log2": round(pmi_log2, 4) if pmi_log2 is not None else None,
            "lookback_steps": PMI_LOOKBACK,
        }
    return out


def write_summary_md(agg: dict, pmi: dict, out_path: Path) -> None:
    cells = agg["cells"]
    arms = sorted({c["arm"] for c in cells})
    scenarios = sorted({c["scenario"] for c in cells})

    lines = ["# OpsGuard memory ablation\n"]
    lines.append("4-arm ablation showing the causal contribution of memory.\n")
    lines.append("## Arms\n")
    for k, v in ARMS.items():
        lines.append(f"- **{k}** — {v['description']}")
    lines.append("")

    for s in scenarios:
        lines.append(f"\n## {s}\n")
        lines.append("| Arm | reward (mean ± std) | attack_recall | spam_caught / total |")
        lines.append("|---|---:|---:|---:|")
        for a in arms:
            cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None)
            if not cell:
                continue
            ar = (
                f"{cell['attack_recall_mean']:.2f}"
                if cell["attack_recall_mean"] is not None else "n/a"
            )
            lines.append(
                f"| {a} | {cell['reward_mean']:+.2f} ± {cell['reward_std']:.2f} | "
                f"{ar} | {cell['spam_caught_mean']:.1f} / {cell['spam_total_mean']:.1f} |"
            )

    # Δ-vs-B table — the headline number.
    lines.append("\n## Reward delta vs Arm B (memory disabled), per scenario\n")
    lines.append("| Scenario | A − B | C − B | D − B |")
    lines.append("|---|---:|---:|---:|")
    for s in scenarios:
        b_cell = next((c for c in cells if c["arm"] == "B_memory_disabled" and c["scenario"] == s), None)
        b_r = b_cell["reward_mean"] if b_cell else 0.0
        deltas = []
        for a in ("A_no_memory_query", "C_memory_aware", "D_memory_aware_profile"):
            cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None)
            if cell is None:
                deltas.append("n/a")
            else:
                d = cell["reward_mean"] - b_r
                deltas.append(f"{d:+.2f}")
        lines.append(f"| {s} | {deltas[0]} | {deltas[1]} | {deltas[2]} |")

    lines.append("\n## PMI: P(catch | recent memory hit) / P(catch), per arm\n")
    lines.append(f"Lookback = {PMI_LOOKBACK} steps. Values > 1.0 ⇒ memory retrieval lifts catch probability.\n")
    lines.append("| Arm | n_catch | n_caught | n_with_recent_hit | P(catch) | P(catch|hit) | PMI (ratio) | PMI (log2) |")
    lines.append("|---|---:|---:|---:|---:|---:|---:|---:|")
    for a in arms:
        p = pmi.get(a, {})
        ratio = p.get("pmi_ratio_p_catch_given_retrieved_over_p_catch")
        log2 = p.get("pmi_log2")
        ratio_s = f"{ratio:.3f}" if ratio is not None else "n/a"
        log2_s = f"{log2:+.3f}" if log2 is not None else "n/a"
        lines.append(
            f"| {a} | {p.get('n_catch_events', 0)} | {p.get('n_caught', 0)} | "
            f"{p.get('n_with_recent_memory_hit', 0)} | "
            f"{p.get('p_catch', 0):.3f} | {p.get('p_catch_given_retrieved', 0):.3f} | "
            f"{ratio_s} | {log2_s} |"
        )

    out_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def plot_delta_vs_b(agg: dict, out_path: Path) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    cells = agg["cells"]
    scenarios = sorted({c["scenario"] for c in cells})
    arms_to_show = ["A_no_memory_query", "C_memory_aware", "D_memory_aware_profile"]

    deltas: dict[str, list[float]] = {a: [] for a in arms_to_show}
    for s in scenarios:
        b_cell = next((c for c in cells if c["arm"] == "B_memory_disabled" and c["scenario"] == s), None)
        b_r = b_cell["reward_mean"] if b_cell else 0.0
        for a in arms_to_show:
            cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None)
            deltas[a].append((cell["reward_mean"] - b_r) if cell else 0.0)

    fig, ax = plt.subplots(figsize=(10, 5))
    width = 0.8 / len(arms_to_show)
    for i, a in enumerate(arms_to_show):
        x = [j + i * width - 0.4 + width / 2 for j in range(len(scenarios))]
        ax.bar(x, deltas[a], width=width * 0.95, label=a)
    ax.set_xticks(range(len(scenarios)))
    ax.set_xticklabels(scenarios, rotation=15, ha="right")
    ax.set_ylabel("Reward delta vs Arm B (memory disabled)")
    ax.set_title("OpsGuard memory ablation: reward lift over no-memory baseline")
    ax.axhline(0, color="gray", linestyle="--", linewidth=0.8)
    ax.legend(loc="best", fontsize=9)
    ax.grid(axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(out_path, dpi=130)
    plt.close(fig)


# ------------------------------------------------------------------------------
# Main.
# ------------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scenarios", nargs="*", default=DEFAULT_SCENARIOS)
    ap.add_argument("--seeds", nargs="*", type=int, default=DEFAULT_SEEDS)
    ap.add_argument("--out", default="eval_outputs/memory_ablation")
    args = ap.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    rollouts_path = out_dir / "rollouts.jsonl"

    records: list[dict] = []
    saved_env: dict[str, str | None] = {}

    print(
        f"[memory_ablation] arms={list(ARMS)} scenarios={args.scenarios} "
        f"seeds={args.seeds} -> {out_dir}"
    )
    t_total = time.time()
    with open(rollouts_path, "w", encoding="utf-8") as fout:
        for arm_name, spec in ARMS.items():
            # Apply env overrides for this arm; restore on exit.
            for k, v in spec["env_overrides"].items():
                saved_env[k] = os.environ.get(k)
                os.environ[k] = v
            try:
                for sid in args.scenarios:
                    for seed in args.seeds:
                        env = OpsguardEnvironment()
                        policy_fn = spec["policy_factory"]()
                        rec = rollout_with_trace(env, arm_name, policy_fn, sid, seed)
                        records.append(rec)
                        fout.write(json.dumps(rec) + "\n")
                        fout.flush()
                        recall = (
                            rec["n_spam_caught"] / rec["n_spam_total"]
                            if rec["n_spam_total"] else float("nan")
                        )
                        print(
                            f"  {arm_name:>26} | {sid:<26} seed={seed} | "
                            f"reward={rec['cumulative_reward']:>+8.2f} "
                            f"recall={recall:.2f} "
                            f"steps={rec['n_steps']:>3} "
                            f"({rec['elapsed_sec']:.1f}s)",
                            flush=True,
                        )
            finally:
                # Restore env vars.
                for k in spec["env_overrides"]:
                    prev = saved_env.get(k)
                    if prev is None:
                        os.environ.pop(k, None)
                    else:
                        os.environ[k] = prev

    agg = aggregate_records(records)
    pmi = compute_pmi(records)

    (out_dir / "summary.json").write_text(json.dumps(agg, indent=2), encoding="utf-8")
    (out_dir / "pmi.json").write_text(json.dumps(pmi, indent=2), encoding="utf-8")
    write_summary_md(agg, pmi, out_dir / "summary.md")
    plot_delta_vs_b(agg, out_dir / "delta.png")

    print(
        f"\n[done] {len(records)} rollouts in {time.time() - t_total:.1f}s -> {out_dir}"
    )
    print(f"  rollouts: {rollouts_path}")
    print(f"  summary:  {out_dir / 'summary.md'}")
    print(f"  pmi:      {out_dir / 'pmi.json'}")
    print(f"  plot:     {out_dir / 'delta.png'}")


if __name__ == "__main__":
    main()