#!/usr/bin/env python3
"""
QA4PC ablation: does the formal model layer improve yes/no/maybe accuracy?

Conditions
----------
  direct   policy + scenario + question → LLM → yes/no/maybe
  formal   policy → SpecificationAnalyzerAgent blueprint →
             blueprint + policy + scenario → LLM → yes/no/maybe

The 133 unique policy trees are processed once and their blueprints cached on
disk (``artifacts/qa4pc_cache/blueprints/``), so re-runs with different sample
sizes don't re-hit the API.

Usage
-----
  python scripts/eval/eval_qa4pc_ablation.py
  python scripts/eval/eval_qa4pc_ablation.py --n 50 --seed 99
  python scripts/eval/eval_qa4pc_ablation.py --build-cache-only
  python scripts/eval/eval_qa4pc_ablation.py --conditions direct   # skip formal
  python scripts/eval/eval_qa4pc_ablation.py --json-out artifacts/qa4pc_ablation.json
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import time
from collections import Counter
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))

from dotenv import load_dotenv
load_dotenv(ROOT / ".env", override=True)

os.environ.setdefault("LLM_BACKEND_FALLBACK", "openai")
os.environ.setdefault("RAG_FALLBACK_MODEL_NAME", "gpt-4o-mini")

from frame.rag_component.llm import LLM
from frame.timed_automata.nl2formalmodel.specification_analyzer import SpecificationAnalyzerAgent

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

CACHE_DIR = ROOT / "artifacts" / "qa4pc_cache" / "blueprints"


# ---------------------------------------------------------------------------
# Dataset loading
# ---------------------------------------------------------------------------

def load_qa4pc_test() -> list[dict]:
    """Download QA4PC test_entailment split and return as list of dicts."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        print("[error] huggingface_hub not installed; run: pip install huggingface_hub")
        sys.exit(1)

    path = hf_hub_download(
        "Marzipan/QA4PC",
        "test_entailment_qa4pc.json",
        repo_type="dataset",
    )
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def sample_items(items: list[dict], n: int, seed: int) -> list[dict]:
    import random
    rng = random.Random(seed)
    if n >= len(items):
        return list(items)
    return rng.sample(items, n)


# ---------------------------------------------------------------------------
# LLM prompts
# ---------------------------------------------------------------------------

_ANSWER_SYSTEM = (
    "You are a policy compliance evaluator. "
    "Given a policy excerpt, a user scenario, and a question, "
    "decide whether the answer is **yes**, **no**, or **maybe** "
    "(maybe = cannot be determined from the policy alone). "
    "Reply with a single word: yes, no, or maybe. No explanation."
)

_DIRECT_TEMPLATE = """\
## Policy
{policy}

## User scenario
{scenario}

## Question
{question}

Reply with exactly one word: yes, no, or maybe."""

_FORMAL_TEMPLATE = """\
## Policy
{policy}

## Formal model of the policy (structured blueprint)
{blueprint}

## User scenario
{scenario}

## Question
{question}

Reply with exactly one word: yes, no, or maybe."""


def parse_ynm(raw: str) -> str | None:
    """Extract yes / no / maybe from LLM output."""
    s = (raw or "").strip().lower()
    for word in re.split(r"[\s.,;:!?]+", s):
        if word in ("yes", "no", "maybe"):
            return word
    return None


# ---------------------------------------------------------------------------
# Blueprint cache
# ---------------------------------------------------------------------------

def blueprint_path(tree_id: str) -> Path:
    return CACHE_DIR / f"{tree_id}.json"


def load_blueprint_cache(tree_id: str) -> dict | None:
    p = blueprint_path(tree_id)
    if p.exists():
        with open(p, encoding="utf-8") as f:
            return json.load(f)
    return None


def save_blueprint_cache(tree_id: str, data: dict) -> None:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    with open(blueprint_path(tree_id), "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def build_blueprints(
    items: list[dict],
    *,
    model_name: str,
    sleep_s: float,
    verbose: bool,
) -> dict[str, dict]:
    """Build (or load from cache) the formal blueprint for every unique tree."""
    trees: dict[str, str] = {}
    for item in items:
        tid = item["tree_id"]
        if tid not in trees:
            trees[tid] = item["policy"]

    analyzer = SpecificationAnalyzerAgent(model_name=model_name)
    blueprints: dict[str, dict] = {}

    for i, (tid, policy) in enumerate(trees.items()):
        cached = load_blueprint_cache(tid)
        if cached is not None:
            blueprints[tid] = cached
            if verbose:
                print(f"  [cache] {tid[:12]}…")
            continue

        if verbose:
            print(f"  [build {i+1}/{len(trees)}] {tid[:12]}…")

        try:
            result = analyzer.analyze(policy)
            data = {
                "prose": result.prose,
                "blueprint": result.blueprint,
            }
        except Exception as exc:
            print(f"  [warn] Analyzer failed for {tid[:12]}: {exc}")
            data = {"prose": "", "blueprint": {}}

        save_blueprint_cache(tid, data)
        blueprints[tid] = data

        if sleep_s > 0:
            time.sleep(sleep_s)

    return blueprints


def _fmt_blueprint(bp: dict) -> str:
    """Compact text representation of the §8 blueprint JSON."""
    if not bp:
        return "(no structured blueprint available)"
    return json.dumps(bp, ensure_ascii=False, indent=2)


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def run_condition(
    items: list[dict],
    *,
    condition: str,
    blueprints: dict[str, dict],
    model_name: str,
    sleep_s: float,
    verbose: bool,
) -> list[dict]:
    llm = LLM(_ANSWER_SYSTEM, model_name=model_name)
    results: list[dict] = []

    for idx, item in enumerate(items):
        policy = item["policy"]
        scenario = item["scenario"]
        question = item["question"]
        gt = item["answer"]  # yes / no / maybe

        if condition == "direct":
            prompt = _DIRECT_TEMPLATE.format(
                policy=policy,
                scenario=scenario,
                question=question,
            )
        else:  # formal
            bp_data = blueprints.get(item["tree_id"], {})
            bp_json = _fmt_blueprint(bp_data.get("blueprint", {}))
            prompt = _FORMAL_TEMPLATE.format(
                policy=policy,
                blueprint=bp_json,
                scenario=scenario,
                question=question,
            )

        raw = llm.generate(user_prompt=prompt)
        pred = parse_ynm(raw)
        correct = (pred == gt) if pred else False

        row = {
            "tree_id": item["tree_id"],
            "utterance_id": item.get("utterance_id", ""),
            "gt": gt,
            "pred": pred or "?",
            "raw": raw.strip()[:120],
            "correct": correct,
        }
        results.append(row)

        if verbose:
            mark = "✓" if correct else "✗"
            print(f"  [{idx+1:3d}/{len(items)}] {mark}  gt={gt:<5}  pred={pred}")

        if sleep_s > 0:
            time.sleep(sleep_s)

    return results


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def compute_metrics(results: list[dict]) -> dict:
    total = len(results)
    if total == 0:
        return {}
    correct = sum(1 for r in results if r["correct"])
    acc = correct / total

    # Per-label breakdown
    by_gt: dict[str, list[bool]] = {}
    for r in results:
        gt = r["gt"]
        by_gt.setdefault(gt, []).append(r["correct"])
    per_label = {lbl: sum(hits) / len(hits) for lbl, hits in by_gt.items()}

    # Confusion
    gt_dist = Counter(r["gt"] for r in results)
    pred_dist = Counter(r["pred"] for r in results)

    return {
        "n": total,
        "accuracy": round(acc, 4),
        "correct": correct,
        "per_label_accuracy": {k: round(v, 4) for k, v in sorted(per_label.items())},
        "gt_distribution": dict(gt_dist),
        "pred_distribution": dict(pred_dist),
    }


def print_report(condition: str, metrics: dict) -> None:
    print(f"\n{'─'*50}")
    print(f"Condition: {condition.upper()}")
    print(f"  n={metrics['n']}  accuracy={metrics['accuracy']:.1%}  correct={metrics['correct']}")
    print(f"  Per-label: {metrics['per_label_accuracy']}")
    print(f"  GT  dist:  {metrics['gt_distribution']}")
    print(f"  Pred dist: {metrics['pred_distribution']}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description="QA4PC ablation: direct vs. formal layer")
    ap.add_argument("--n", type=int, default=100, help="items to evaluate (default: 100)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument(
        "--conditions",
        nargs="+",
        choices=["direct", "formal"],
        default=["direct", "formal"],
    )
    ap.add_argument("--model", default="gpt-4o-mini", help="LLM model for answering (default: gpt-4o-mini)")
    ap.add_argument("--analyzer-model", default="gpt-4.1", help="LLM model for Analyzer agent")
    ap.add_argument("--sleep", type=float, default=0.5, help="seconds between API calls")
    ap.add_argument("--build-cache-only", action="store_true", help="only build blueprint cache, no eval")
    ap.add_argument("--json-out", type=Path, default=None)
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()

    print("Loading QA4PC test_entailment split…")
    all_items = load_qa4pc_test()
    print(f"  Loaded {len(all_items)} items ({len({x['tree_id'] for x in all_items})} unique trees)")

    items = sample_items(all_items, args.n, args.seed)
    print(f"  Sampled {len(items)} items (seed={args.seed})")

    blueprints: dict[str, dict] = {}
    if "formal" in args.conditions or args.build_cache_only:
        print(f"\nBuilding/loading formal blueprints (model={args.analyzer_model})…")
        blueprints = build_blueprints(
            items,
            model_name=args.analyzer_model,
            sleep_s=args.sleep,
            verbose=args.verbose,
        )
        cached_count = sum(
            1 for item in items
            if blueprint_path(item["tree_id"]).exists()
        )
        print(f"  Done — {cached_count}/{len({x['tree_id'] for x in items})} trees cached")

    if args.build_cache_only:
        print("--build-cache-only: stopping after cache build.")
        return

    all_results: dict[str, Any] = {"conditions": {}}

    for cond in args.conditions:
        print(f"\nRunning condition: {cond.upper()} (model={args.model})…")
        results = run_condition(
            items,
            condition=cond,
            blueprints=blueprints,
            model_name=args.model,
            sleep_s=args.sleep,
            verbose=args.verbose,
        )
        metrics = compute_metrics(results)
        all_results["conditions"][cond] = {"metrics": metrics, "rows": results}
        print_report(cond, metrics)

    # Summary comparison
    if len(args.conditions) > 1:
        print(f"\n{'═'*50}")
        print("Summary")
        for cond in args.conditions:
            m = all_results["conditions"][cond]["metrics"]
            print(f"  {cond:<8}  acc={m['accuracy']:.1%}  ({m['correct']}/{m['n']})")
        conds = args.conditions
        if len(conds) == 2:
            a0 = all_results["conditions"][conds[0]]["metrics"]["accuracy"]
            a1 = all_results["conditions"][conds[1]]["metrics"]["accuracy"]
            delta = a1 - a0
            print(f"\n  Δ ({conds[1]} − {conds[0]}) = {delta:+.1%}")

    if args.json_out:
        args.json_out.parent.mkdir(parents=True, exist_ok=True)
        all_results["config"] = vars(args)
        with open(args.json_out, "w", encoding="utf-8") as f:
            json.dump(all_results, f, ensure_ascii=False, indent=2, default=str)
        print(f"\nResults saved → {args.json_out}")


if __name__ == "__main__":
    main()