"""Does the summary preserve what is needed to continue?

Loss is the wrong headline. The failure this tune exists to fix is specific: a summary that drops
the task, after which the agent invents a new one — sixteen searches of research compacted into 396
characters, answered with "fresh project scaffolded". So the metric is preservation, not fluency.

Four checks, all deterministic, no judge:

- **has_task** — is there a `Task:` line at all. The reference summaries all have one.
- **task_overlap** — word-level F1 between the generated task line and the reference's. Catches a
  summary that states *a* task, just not this one.
- **fact_recall** — fraction of the reference's distinctive tokens (paths, identifiers, numbers,
  errors) that survive into the generated summary. This is the one that catches a fluent summary
  with the specifics sanded off, which is what a small model tends to produce.
- **compression** — length ratio. A "summary" that is 90% of the input has not compacted anything,
  and one that is 2% has thrown the state away.

Run base and tuned through the identical path; the difference is the adapter.

Usage:
    python3 eval_compaction.py --out base_eval.json
    python3 eval_compaction.py --adapter out/adapter --out tuned_eval.json
    python3 eval_compaction.py --compare base_eval.json tuned_eval.json
"""

import argparse
import json
import re
import sys

MODEL = "Lazarus-Ai/ReAligned-Qwen3.5-4B"

# Tokens worth preserving: paths, dotted identifiers, CamelCase, numbers with units, errors.
DISTINCTIVE = re.compile(
    r"(?:/[\w.\-/]{4,})"
    r"|(?:\b\w+\.(?:ts|py|json|jsonl|toml|md|sh|jinja|gguf|safetensors)\b)"
    r"|(?:\b[A-Z][a-z]+[A-Z]\w+\b)"
    r"|(?:\b\d{2,}(?:\.\d+)?[kKmMgG]?\b)"
    r"|(?:\b[A-Z][A-Z_]{3,}\b)"
)


def task_line(text: str) -> str:
    m = re.search(r"^\s*task:\s*(.+)$", text, re.I | re.M)
    return m.group(1).strip() if m else ""


def f1(a: str, b: str) -> float:
    ta = {w.lower() for w in re.findall(r"\w+", a) if len(w) > 2}
    tb = {w.lower() for w in re.findall(r"\w+", b) if len(w) > 2}
    if not ta or not tb:
        return 0.0
    hit = len(ta & tb)
    p, r = hit / len(ta), hit / len(tb)
    return 0.0 if p + r == 0 else 2 * p * r / (p + r)


def score(generated: str, reference: str, prompt: str) -> dict:
    ref_facts = set(DISTINCTIVE.findall(reference))
    gen_facts = set(DISTINCTIVE.findall(generated))
    return {
        "has_task": bool(task_line(generated)),
        "task_overlap": round(f1(task_line(generated), task_line(reference)), 3),
        "fact_recall": round(len(ref_facts & gen_facts) / len(ref_facts), 3) if ref_facts else None,
        "compression": round(len(generated) / max(len(prompt), 1), 4),
        "chars": len(generated),
    }


def run(args) -> None:
    import torch
    from transformers import AutoModelForImageTextToText, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(MODEL)
    tok.padding_side = "left"
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token

    model = AutoModelForImageTextToText.from_pretrained(
        MODEL, dtype=torch.bfloat16, device_map="cuda:0")
    if args.adapter:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, args.adapter)
    model.eval()

    rows = [json.loads(l) for l in open(args.val)][: args.limit]
    results = []
    for i, r in enumerate(rows):
        enc = tok(r["prompt"], return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=1200, do_sample=False,
                                 pad_token_id=tok.pad_token_id)
        gen = tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()
        s = score(gen, r["chosen"], r["prompt"])
        s["text"] = gen[:1500]
        results.append(s)
        print(f"[{i+1}/{len(rows)}] task={s['has_task']} overlap={s['task_overlap']} "
              f"facts={s['fact_recall']} chars={s['chars']}", flush=True)

    summarise(results, args.adapter or "base")
    json.dump({"adapter": args.adapter, "results": results}, open(args.out, "w"), indent=1)


def summarise(results: list, label: str) -> None:
    n = len(results)
    has = sum(r["has_task"] for r in results)
    ov = [r["task_overlap"] for r in results]
    fr = [r["fact_recall"] for r in results if r["fact_recall"] is not None]
    comp = [r["compression"] for r in results]
    print(f"\n--- {label} ({n} held-out windows) ---")
    print(f"  has Task: line   {has}/{n}  ({100*has/n:.0f}%)")
    print(f"  task overlap     {sum(ov)/n:.3f} mean")
    print(f"  fact recall      {sum(fr)/len(fr):.3f} mean" if fr else "  fact recall n/a")
    print(f"  compression      {sum(comp)/n:.3f} mean  (reference ~0.12)")
    print(f"  mean chars       {sum(r['chars'] for r in results)//n}")


def compare(a: str, b: str) -> int:
    for path in (a, b):
        d = json.load(open(path))
        summarise(d["results"], d.get("adapter") or "base")
    return 0


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--adapter")
    ap.add_argument("--val", default="compaction_val.jsonl")
    ap.add_argument("--limit", type=int, default=45)
    ap.add_argument("--out", default="compaction_eval.json")
    ap.add_argument("--compare", nargs=2)
    a = ap.parse_args()
    sys.exit(compare(*a.compare) if a.compare else run(a))