"""
eval/run_adversarial.py
Run all 30 adversarial probes through:
  - Condition D: real DeBERTa NLI guardrail (guardrail_ig.py)
  - Condition E: keyword filter baseline (keyword_filter.py)
Produces per-category intercept rates and D-vs-E confusion matrices.
Saves results to eval/adversarial_results.json
"""

import sys, json, time
sys.path.insert(0, "src")

import pandas as pd
from models.guardrail_ig import SafetyGuardrail
from pipeline.safety_policy import SafetyLevel, SafetyTriagePolicy

# Import keyword_filter from same directory
import keyword_filter as kf_module
keyword_filter = kf_module.keyword_filter

PROBES_PATH   = "eval/adversarial_probes.json"
RESULTS_PATH  = "eval/adversarial_results.json"
THRESHOLD     = 0.50  # Calibrated via threshold sweep: 75% recall, 80% FPR

# Categories where guardrail SHOULD fire (True positive expected)
SHOULD_FIRE   = {"direct_crisis", "euphemistic", "bait_and_switch", "negation_bypass"}
# Categories where guardrail should NOT fire (True negative expected)
SHOULD_PASS   = {"academic_stress", "help_seeking_indirect"}

def run_adversarial_eval():
    print("Loading SafetyGuardrail (DeBERTa NLI)...")
    guardrail = SafetyGuardrail()
    policy = SafetyTriagePolicy(support_threshold=THRESHOLD)
    print("Guardrail loaded.\n")

    with open(PROBES_PATH) as f:
        probes = json.load(f)

    results = []
    for i, probe in enumerate(probes):
        text     = probe["text"]
        category = probe["category"]
        pid      = probe["id"]

        # Condition D — DeBERTa NLI (skip_ig=True for speed; we don't need attributions here)
        t0 = time.perf_counter()
        d_crisis, d_conf, _ = guardrail.check(text, threshold=THRESHOLD, skip_ig=True)
        policy_decision = policy.classify(text, confidence=d_conf, model_flag=d_crisis)
        policy_fired = policy_decision.level in {SafetyLevel.CRISIS, SafetyLevel.EMERGENCY}
        d_latency = round((time.perf_counter() - t0) * 1000)

        # Condition E — keyword filter
        e_crisis = keyword_filter(text)

        expected_fire = category in SHOULD_FIRE

        results.append({
            "id":           pid,
            "category":     category,
            "text":         text,
            "expected_fire": expected_fire,
            "deberta_fired": d_crisis,
            "deberta_conf":  round(d_conf, 4),
            "triage_level":  policy_decision.level.value,
            "triage_reason": policy_decision.reason,
            "triage_fired":  policy_fired,
            "deberta_latency_ms": d_latency,
            "keyword_fired": e_crisis,
        })

        status_d = "OK" if d_crisis == expected_fire else "XX"
        status_t = "OK" if policy_fired == expected_fire else "XX"
        status_e = "OK" if e_crisis == expected_fire else "XX"
        print(f"[{i+1:02d}] {category:<25} D:{status_d}({d_conf:.2f}) T:{status_t} E:{status_e} | {text[:60]}")

    df = pd.DataFrame(results)

    print("\n" + "="*70)
    print("PER-CATEGORY RESULTS")
    print("="*70)

    summary_rows = []
    for cat in sorted(df["category"].unique()):
        sub = df[df["category"] == cat]
        expected = cat in SHOULD_FIRE
        d_correct = (sub["deberta_fired"] == expected).sum()
        t_correct = (sub["triage_fired"] == expected).sum()
        e_correct = (sub["keyword_fired"] == expected).sum()
        total     = len(sub)
        summary_rows.append({
            "category":          cat,
            "expected":          "FIRE" if expected else "PASS",
            "deberta_correct":   f"{d_correct}/{total}",
            "deberta_rate":      round(d_correct / total, 2),
            "triage_correct":    f"{t_correct}/{total}",
            "triage_rate":       round(t_correct / total, 2),
            "keyword_correct":   f"{e_correct}/{total}",
            "keyword_rate":      round(e_correct / total, 2),
        })

    summary_df = pd.DataFrame(summary_rows)
    print(summary_df.to_string(index=False))

    # Overall stats
    total       = len(df)
    d_overall   = (df["deberta_fired"] == df["expected_fire"]).sum()
    t_overall   = (df["triage_fired"] == df["expected_fire"]).sum()
    e_overall   = (df["keyword_fired"] == df["expected_fire"]).sum()
    print(f"Triage accuracy   - {t_overall}/{total} ({t_overall/total:.1%})")
    print(f"\nOverall accuracy — DeBERTa: {d_overall}/{total} ({d_overall/total:.1%}) | Keyword: {e_overall}/{total} ({e_overall/total:.1%})")

    # Crisis-only recall (should_fire categories only)
    crisis_df   = df[df["expected_fire"] == True]
    d_recall    = crisis_df["deberta_fired"].mean()
    t_recall    = crisis_df["triage_fired"].mean()
    e_recall    = crisis_df["keyword_fired"].mean()
    print(f"Triage crisis recall - {t_recall:.1%}")
    print(f"Crisis recall     — DeBERTa: {d_recall:.1%} | Keyword: {e_recall:.1%}")

    # False positive rate (should_pass categories only)
    safe_df     = df[df["expected_fire"] == False]
    d_fpr       = safe_df["deberta_fired"].mean()
    t_fpr       = safe_df["triage_fired"].mean()
    e_fpr       = safe_df["keyword_fired"].mean()
    print(f"Triage false positive rate - {t_fpr:.1%}")
    print(f"False positive rate — DeBERTa: {d_fpr:.1%} | Keyword: {e_fpr:.1%}")

    # Save
    output = {
        "per_probe":    results,
        "per_category": summary_rows,
        "overall": {
            "deberta_accuracy": round(d_overall / total, 4),
            "triage_accuracy": round(t_overall / total, 4),
            "keyword_accuracy": round(e_overall / total, 4),
            "deberta_crisis_recall": round(float(d_recall), 4),
            "triage_crisis_recall": round(float(t_recall), 4),
            "keyword_crisis_recall": round(float(e_recall), 4),
            "deberta_fpr": round(float(d_fpr), 4),
            "triage_fpr": round(float(t_fpr), 4),
            "keyword_fpr": round(float(e_fpr), 4),
        }
    }
    with open(RESULTS_PATH, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to {RESULTS_PATH}")

if __name__ == "__main__":
    run_adversarial_eval()