src/training/run_ablations.py

"""
Ablation study — evaluates the three controllable ablations on the validation set.

A1  No Tool Context  — [ACTION]+[OBS] tokens zeroed at preprocessing time
A2  No Query Context — [QUERY] tokens zeroed at preprocessing time
A3  No Cross-Step    — the TransformerEncoder is bypassed (each step independent)

All ablations reuse the *same trained weights* — they are forward-pass
modifications, not retrains.  This is correct because we are measuring
the sensitivity of the FULL model to each input channel.

Output: val_ablations.json
"""
import os
import sys
import json
import torch

script_dir   = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.join(script_dir, "..", "..")
sys.path.insert(0, project_root)

from src.models.agentsight import AgentSightModel
from src.data.preprocessor  import StepPreprocessor
from src.training.evaluate  import step_localization_accuracy
from sklearn.metrics import f1_score


# ── Ablated preprocessors ─────────────────────────────────────────────────────

class AblatedPreprocessor(StepPreprocessor):
    """Zeroes out one or more input channels before encoding."""
    def __init__(self, ablate_tool=False, ablate_query=False, **kwargs):
        super().__init__(**kwargs)
        self.ablate_tool  = ablate_tool
        self.ablate_query = ablate_query

    def encode_step(self, query, step):
        if self.ablate_tool:
            step = dict(step)
            step["tool_calls"]     = []
            step["tool_responses"] = []
        if self.ablate_query:
            query = ""
        return super().encode_step(query, step)


# ── No-cross-step variant ─────────────────────────────────────────────────────

class NoContextModel(AgentSightModel):
    """Bypass the TransformerEncoder — each step is classified independently."""
    def forward(self, input_ids, attention_mask):
        out       = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        step_repr = out.last_hidden_state[:, 0, :].to(torch.float32)
        fused     = self.fusion(step_repr)
        return self.cls_head(fused).squeeze(-1)


# ── Runner ────────────────────────────────────────────────────────────────────

def run_condition(name, model, preprocessor, val_samples, device, threshold=0.5):
    model.eval()
    results = []

    with torch.no_grad():
        for i, sample in enumerate(val_samples):
            is_hal = sample.get("is_hallucination", False)
            if isinstance(is_hal, str):
                is_hal = is_hal.lower() == "true"
            hal_step = sample.get("hallucination_step")
            if hal_step is not None:
                hal_step = int(hal_step)

            try:
                steps = preprocessor.encode_trajectory(sample)
            except Exception:
                steps = []

            if not steps:
                results.append({
                    "condition": name,
                    "sample_idx": i,
                    "true_is_hallucination": is_hal,
                    "true_hallucination_step": hal_step,
                    "pred_is_hallucination": False,
                    "pred_step": None,
                    "max_hal_prob": None,
                    "encoding_failed": True,
                })
                continue

            ids  = torch.stack([s["encoding"]["input_ids"].squeeze(0)  for s in steps]).to(device)
            mask = torch.stack([s["encoding"]["attention_mask"].squeeze(0) for s in steps]).to(device)
            vocab_size = model.encoder.config.vocab_size
            ids = torch.clamp(ids, 0, vocab_size - 1)

            logits     = model(ids, mask)
            probs      = torch.sigmoid(logits).cpu().tolist()
            if isinstance(probs, float):
                probs = [probs]

            max_prob    = max(probs)
            pred_is_hal = max_prob > threshold
            pred_step   = steps[probs.index(max_prob)]["step_idx"] if pred_is_hal else None

            results.append({
                "condition": name,
                "sample_idx": i,
                "true_is_hallucination": is_hal,
                "true_hallucination_step": hal_step,
                "pred_is_hallucination": pred_is_hal,
                "pred_step": pred_step,
                "max_hal_prob": max_prob,
                "encoding_failed": False,
            })

    return results


def compute_metrics(results, val_samples):
    hal_true  = [1 if r["true_is_hallucination"] else 0 for r in results]
    hal_preds = [1 if r["pred_is_hallucination"] else 0 for r in results]
    f1 = f1_score(hal_true, hal_preds, average="macro", zero_division=0)

    # Step-acc: official denominator = all hallucinated, not TP-only
    hal_only = [r for r in results if r["true_is_hallucination"] and not r.get("encoding_failed")]
    correct  = sum(1 for r in hal_only if r["pred_step"] == r["true_hallucination_step"])
    step_acc = correct / len(hal_only) if hal_only else 0.0

    return {"step_acc": step_acc, "judgment_f1": f1,
            "step_correct": correct, "n_hal": len(hal_only)}


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    weights_path = os.path.join(project_root, "src", "models", "best_agentsight.pth")
    meta_path    = weights_path.replace(".pth", "_meta.json")

    threshold = 0.5
    if os.path.exists(meta_path):
        with open(meta_path) as f:
            meta = json.load(f)
        threshold = meta.get("threshold", 0.5)
        print(f"Using saved threshold: {threshold:.3f}")

    with open(os.path.join(project_root, "data", "splits", "val.json")) as f:
        val_samples = json.load(f)
    print(f"Loaded {len(val_samples)} val samples.\n")

    conditions = [
        ("Full Model",       AgentSightModel, StepPreprocessor()),
        ("No Tool Context",  AgentSightModel, AblatedPreprocessor(ablate_tool=True)),
        ("No Query Context", AgentSightModel, AblatedPreprocessor(ablate_query=True)),
        ("No Cross-Step",    NoContextModel,  StepPreprocessor()),
    ]

    all_results = []
    summary     = {}

    for name, ModelClass, preprocessor in conditions:
        print(f"Running ablation: {name} …")
        model = ModelClass()
        model.load_state_dict(torch.load(weights_path, map_location=device))
        model.to(device)

        results = run_condition(name, model, preprocessor, val_samples, device, threshold)
        all_results.extend(results)

        m = compute_metrics(results, val_samples)
        summary[name] = m
        print(f"  Step-Acc: {m['step_acc']*100:.1f}% ({m['step_correct']}/{m['n_hal']})  "
              f"|  F1: {m['judgment_f1']*100:.1f}%\n")

    # ── Save ──────────────────────────────────────────────────────────────────
    out_path = os.path.join(project_root, "val_ablations.json")
    with open(out_path, "w") as f:
        json.dump({
            "note": "Ablations evaluated on val.json only. test.json is sealed.",
            "test_json_sha256": "9604aae8eb5aec4ae666cfbe3053910f0570a807a4fa5515223dbca1aa66a7d8",
            "threshold_used": threshold,
            "summary": summary,
            "per_trajectory": all_results,
        }, f, indent=2)

    # ── Print summary table ───────────────────────────────────────────────────
    base_acc = summary["Full Model"]["step_acc"]
    print("=" * 64)
    print(f"{'Condition':<24}  {'Step-Acc':>9}  {'Δ vs Full':>9}  {'Macro-F1':>9}")
    print("-" * 64)
    for name, m in summary.items():
        delta = m["step_acc"] - base_acc
        delta_str = f"({delta*100:+.1f})" if name != "Full Model" else "baseline"
        print(f"  {name:<22}  {m['step_acc']*100:>8.1f}%  {delta_str:>9}  "
              f"{m['judgment_f1']*100:>8.1f}%")
    print("=" * 64)
    print(f"\nRaw output → {out_path}")


if __name__ == "__main__":
    main()