"""
Qualitative diff across arms: actually READ the stories.

Aggregate metrics can say "effective rank 2.0/16" without conveying that six of
the sixteen ships are literally named *Aethel*. This module extracts the
concrete signature of collapse -- shared openings, repeated proper nouns,
shared closing cadence, tone/genre spread -- so a human can see what changed.

Works on either a scored pool (`--pool`) or an eval dump (`--eval`), so the
base model, every checkpoint and every arm are read the same way.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent

OPEN_FRAME = re.compile(r"^The\s+\*?([A-Z][\w' ]{2,20})\*?\s+(\w+ed|\w+s)\s+"
                        r"(through|over|above|at|into|across)\b")
ITALIC_NAME = re.compile(r"\*([A-Z][\w' ]{2,24})\*")
CADENCE = [
    (r"for the first time in (centuries|years|generations|decades|a long time)",
     "'for the first time in X'"),
    (r"\b(had )?begun to\b", "'begun to'"),
    (r"\bwasn'?t\s+\w+[^.]{0,50},?\s+it\s+was\b", "'it wasn't X, it was Y'"),
    (r"\n\n[^\n]{1,60}\.\s*$", "one-line closing paragraph"),
]
# crude genre/register probes -- presence of ANY is weak, but the SPREAD across
# a 16-sample set is the interesting number
REGISTER = {
    "elegiac/literary": r"\b(silence|ache|hollow|dust|memory|grief|quiet)\b",
    "comic": r"\b(ridiculous|absurd|joke|laughed|idiot|bureaucra|paperwork)\b",
    "horror": r"\b(scream|blood|rot|teeth|corpse|terror|wrong)\b",
    "technical/SF": r"\b(protocol|sensor|reactor|alloy|deploy|calibrat|drone)\b",
    "dialogue-driven": r'"[^"]{10,}"',
    "second person": r"\bYou (are|were|walk|feel|stand|know)\b",
    "epistolary": r"\b(Dear |entry |log:|Report |memo)\b",
}


def first_sentence(t: str, n: int = 130) -> str:
    return re.split(r"(?<=[.!?])\s", t.strip())[0][:n]


def last_sentence(t: str, n: int = 120) -> str:
    return re.split(r"(?<=[.!?])\s", t.strip())[-1][:n]


def analyze_group(texts: list[str]) -> dict:
    n = len(texts)
    names = Counter(x for t in texts for x in ITALIC_NAME.findall(t))
    opens_the = sum(1 for t in texts if t.strip().startswith("The "))
    frame = sum(1 for t in texts if OPEN_FRAME.match(t.strip()))
    # repeated content words across stories (excluding prompt-driven ones is
    # impossible in general, so we report the top shared nouns as-is)
    stop = set("the a an and or but of to in on it its was were is are that this "
               "with for as at by from had have has been be he she they we you i "
               "his her their our my not no so then than there here what when "
               "which who all more most into over under out up down".split())
    per_story_vocab = [set(w for w in re.sub(r"[^a-z\s]", " ", t.lower()).split()
                           if w not in stop and len(w) > 3) for t in texts]
    shared = Counter(w for v in per_story_vocab for w in v)
    ubiquitous = {w: c for w, c in shared.most_common(400) if c >= max(3, int(0.6 * n))}
    cad = {}
    for pat, label in CADENCE:
        cad[label] = sum(1 for t in texts if re.search(pat, t, re.I | re.M))
    reg = {}
    for label, pat in REGISTER.items():
        reg[label] = sum(1 for t in texts if re.search(pat, t, re.I))
    return {
        "n": n,
        "opens_with_The": opens_the,
        "opening_frame_match": frame,
        "distinct_first_5_words": len({" ".join(t.strip().split()[:5]).lower() for t in texts}),
        "top_names": dict(names.most_common(6)),
        "max_name_reuse": max(names.values()) if names else 0,
        "ubiquitous_words": dict(list(ubiquitous.items())[:12]),
        "cadence": cad,
        "register_spread": reg,
        "registers_present": sum(1 for v in reg.values() if v >= max(2, int(0.15 * n))),
    }


def load_pool(tag: str, split: str, prompt_id: str | None):
    rows = [json.loads(l) for l in
            open(ROOT / "outputs" / f"pool_{tag}" / f"pool_{split}.jsonl") if l.strip()]
    by = defaultdict(list)
    for r in rows:
        by[r["prompt_id"]].append(r)
    if prompt_id:
        return {prompt_id: by[prompt_id]}
    return dict(by)


def load_eval(label: str, out: str):
    p = ROOT / out / "samples" / f"{label}_full.json"
    data = json.load(open(p))
    return {d["prompt_id"]: [{"text": t, "prompt": d["prompt"]} for t in d["texts"]]
            for d in data}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pool", help="pool tag, e.g. 4b")
    ap.add_argument("--eval", help="eval label, e.g. E0-baseline")
    ap.add_argument("--out", default="outputs/eval")
    ap.add_argument("--split", default="train")
    ap.add_argument("--prompt-id", default=None)
    ap.add_argument("--show", type=int, default=16, help="first sentences to print")
    ap.add_argument("--limit-prompts", type=int, default=40)
    args = ap.parse_args()

    groups = load_pool(args.pool, args.split, args.prompt_id) if args.pool \
        else load_eval(args.eval, args.out)
    keys = list(groups)[: args.limit_prompts]

    agg = defaultdict(list)
    for pid in keys:
        texts = [r["text"] for r in groups[pid]]
        if len(texts) < 4:
            continue
        a = analyze_group(texts)
        agg["opens_with_The"].append(a["opens_with_The"] / a["n"])
        agg["opening_frame"].append(a["opening_frame_match"] / a["n"])
        agg["distinct_openers"].append(a["distinct_first_5_words"] / a["n"])
        agg["max_name_reuse"].append(a["max_name_reuse"])
        agg["registers_present"].append(a["registers_present"])
        for k, v in a["cadence"].items():
            agg[f"cadence:{k}"].append(v / a["n"])

    name = args.eval or f"pool-{args.pool}"
    print(f"\n===== QUALITATIVE PROFILE: {name}  ({len(keys)} prompts) =====")
    for k, v in agg.items():
        print(f"  {k:38} {np.mean(v):.3f}")

    pid = args.prompt_id or keys[0]
    texts = [r["text"] for r in groups[pid]][: args.show]
    print(f"\n--- prompt {pid} ---")
    print("   ", groups[pid][0].get("prompt", "")[:220])
    print("\nOPENINGS:")
    for i, t in enumerate(texts):
        print(f"  {i+1:2d}. {first_sentence(t)}")
    print("\nCLOSINGS:")
    for i, t in enumerate(texts):
        print(f"  {i+1:2d}. ...{last_sentence(t)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())