#!/usr/bin/env python3
"""WP44 — taxonomy validation: stratified sample + two independent LLM raters (Anthropic SDK, structured outputs) + metrics.
  --draw     : sample from DuckDB → taxonomy/validation/sample.parquet (+ sample.csv)
  --rate     : rate with --model (one rater per run), resumable, batches of 10 works per request → validation/ratings_<model>.jsonl
  --metrics  : compute agreement metrics → validation/metrics.json, disagreements.csv
API key: read from the file given by --env-file (never from the repository); the key is only used here."""
import argparse, json, os, random, time, re, duckdb, yaml, pandas as pd, numpy as np

SCHEMA = {"type": "object", "properties": {"ratings": {"type": "array", "items": {"type": "object", "properties": {
    "work_id": {"type": "integer"}, "categories": {"type": "array", "items": {"type": "string"}, "maxItems": 5},
    "primary": {"type": ["string", "null"]}, "confidence": {"type": "number"}, "rationale": {"type": "string"}},
    "required": ["work_id", "categories", "primary", "confidence", "rationale"], "additionalProperties": False}}},
    "required": ["ratings"], "additionalProperties": False}


def draw(root, seed, per_cat, n_none):
    con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True)
    con.execute(f"CREATE OR REPLACE VIEW work_primary AS SELECT * FROM read_parquet('{root}/data/processed/classification/primary/*.parquet')")
    con.execute(f"CREATE OR REPLACE VIEW work_categories AS SELECT * FROM read_parquet('{root}/data/processed/classification/cats/*.parquet')")
    df = con.execute(f"""
        WITH base AS (SELECT c.id AS work_id, c.title, a.abstract, c.source_name, c.publication_year AS year, p.primary_category,
                             (SELECT list(category_id) FROM work_categories k WHERE k.work_id = c.id) AS rule_categories,
                             hash(c.id || '{seed}') AS h
                      FROM corpus_works c JOIN oa_abstracts a ON a.work_id = c.id LEFT JOIN work_primary p ON p.work_id = c.id WHERE c.eligible),
             ranked AS (SELECT *, row_number() OVER (PARTITION BY primary_category ORDER BY h) rn FROM base)
        SELECT * EXCLUDE (h, rn) FROM ranked WHERE (primary_category IS NOT NULL AND rn <= {per_cat}) OR (primary_category IS NULL AND rn <= {n_none})""").df()
    rng = random.Random(seed); idx = list(range(len(df))); rng.shuffle(idx)
    df["pi_review"] = False; df.loc[df.index[idx[:200]], "pi_review"] = True
    d = os.path.join(root, "taxonomy/validation"); os.makedirs(d, exist_ok=True)
    df.to_parquet(f"{d}/sample.parquet"); df.drop(columns=["abstract"]).to_csv(f"{d}/sample.csv", index=False)
    print("sample", len(df), "works;", df.primary_category.nunique(), "categories;", df.primary_category.isna().sum(), "unassigned")


def rate(root, model, env_file, ontology, batch):
    import anthropic
    key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(env_file)) if l.startswith("ANTHROPIC_API_KEY="))
    client = anthropic.Anthropic(api_key=key, max_retries=5)
    ont = yaml.safe_load(open(ontology))
    defs = "\n".join(f"- {c['id']}: {c['name']} — {c['definition']}" for c in ont["categories"])
    system = ("You are an expert classifier of computer-science research papers for a systematic mapping study. Assign each paper to the "
              "applicable categories of the ontology below (zero to five; multi-label allowed; choose a primary category when at least one applies). "
              "Judge from the title, abstract, venue and year only. Be conservative: assign a category only when the paper's main "
              "contribution clearly belongs to it. If the paper is not computer-science research, return no categories.\n\nOntology:\n" + defs)
    d = os.path.join(root, "taxonomy/validation"); df = pd.read_parquet(f"{d}/sample.parquet")
    out_p = f"{d}/ratings_{model}.jsonl"; done = set()
    if os.path.exists(out_p):
        done = {json.loads(l)["work_id"] for l in open(out_p)}
    todo = df[~df.work_id.isin(done)]; usage = {"in": 0, "out": 0, "calls": 0}
    with open(out_p, "a") as fh:
        for i in range(0, len(todo), batch):
            b = todo.iloc[i:i + batch]
            items = "\n\n".join(f"### work_id {r.work_id}\nTitle: {r.title}\nVenue: {r.source_name or 'n/a'} ({r.year})\nAbstract: {(r.abstract or '')[:2500]}" for r in b.itertuples())
            for attempt in range(4):
                try:
                    resp = client.messages.create(model=model, max_tokens=8000, system=[{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}],
                                                  messages=[{"role": "user", "content": f"Rate the following {len(b)} papers. Return one rating per work_id.\n\n{items}"}],
                                                  output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"})
                    break
                except anthropic.RateLimitError:
                    time.sleep(20 * (attempt + 1))
            if resp.stop_reason == "refusal":
                print("refusal on batch", i); continue
            text = next(bk.text for bk in resp.content if bk.type == "text"); data = json.loads(text)
            usage["in"] += resp.usage.input_tokens + (resp.usage.cache_read_input_tokens or 0); usage["out"] += resp.usage.output_tokens; usage["calls"] += 1
            for r in data["ratings"]:
                r["model"] = model; fh.write(json.dumps(r) + "\n")
            fh.flush()
            if usage["calls"] % 20 == 0:
                print(f"{usage['calls']} calls, {i+len(b)}/{len(todo)} works, tokens in={usage['in']} out={usage['out']}", flush=True)
    json.dump(usage, open(f"{d}/usage_{model}.json", "w"), indent=1); print("done", usage)


def metrics(root, ontology):
    from sklearn.metrics import cohen_kappa_score
    d = os.path.join(root, "taxonomy/validation"); df = pd.read_parquet(f"{d}/sample.parquet")
    cats = [c["id"] for c in yaml.safe_load(open(ontology))["categories"]]
    raters = {}
    for f in sorted(os.listdir(d)):
        if f.startswith("ratings_") and f.endswith(".jsonl"):
            rows = [json.loads(l) for l in open(os.path.join(d, f))]
            raters[f[8:-6]] = {r["work_id"]: r for r in rows}
    ids = [w for w in df.work_id if all(w in r for r in raters.values())]
    rule = {r.work_id: set(r.rule_categories or []) for r in df.itertuples()}
    rule_p = {r.work_id: r.primary_category for r in df.itertuples()}
    M = {"n_rated_by_all": len(ids), "raters": list(raters)}
    def prf(ref_sets, pred_sets):
        out = {}
        for c in cats:
            tp = sum(1 for w in ids if c in ref_sets[w] and c in pred_sets[w]); fp = sum(1 for w in ids if c not in ref_sets[w] and c in pred_sets[w])
            fn = sum(1 for w in ids if c in ref_sets[w] and c not in pred_sets[w])
            p = tp / (tp + fp) if tp + fp else None; r = tp / (tp + fn) if tp + fn else None
            out[c] = {"tp": tp, "fp": fp, "fn": fn, "precision": p, "recall": r, "f1": (2 * p * r / (p + r)) if p and r else (0.0 if (tp + fp + fn) else None)}
        micro_tp = sum(v["tp"] for v in out.values()); micro_fp = sum(v["fp"] for v in out.values()); micro_fn = sum(v["fn"] for v in out.values())
        mp = micro_tp / max(micro_tp + micro_fp, 1); mr = micro_tp / max(micro_tp + micro_fn, 1)
        f1s = [v["f1"] for v in out.values() if v["f1"] is not None]
        return {"per_category": out, "micro_precision": mp, "micro_recall": mr, "micro_f1": 2 * mp * mr / max(mp + mr, 1e-9), "macro_f1": float(np.mean(f1s))}
    for name, R in raters.items():
        sets = {w: set(x for x in R[w]["categories"] if x in cats) for w in ids}
        M[f"rules_vs_{name}"] = prf(sets, rule)
        M[f"kappa_primary_rules_vs_{name}"] = float(cohen_kappa_score([rule_p.get(w) or "none" for w in ids], [R[w]["primary"] or "none" for w in ids]))
    if len(raters) >= 2:
        (a, A), (b, B) = list(raters.items())[:2]
        M[f"kappa_primary_{a}_vs_{b}"] = float(cohen_kappa_score([A[w]["primary"] or "none" for w in ids], [B[w]["primary"] or "none" for w in ids]))
        cons = {w: set(A[w]["categories"]) & set(B[w]["categories"]) for w in ids}
        M["rules_vs_consensus"] = prf(cons, rule)
        # Krippendorff's alpha (interval form on MASI distance between the two raters' label sets; Passonneau 2006)
        def masi(x, y):
            x, y = set(x), set(y)
            if not x and not y: return 0.0
            j = len(x & y) / len(x | y); m = 1.0 if x == y else (0.67 if (x < y or y < x) else (0.33 if x & y else 0.0))
            return 1.0 - j * m
        pairs = [(A[w]["categories"], B[w]["categories"]) for w in ids]
        Do = float(np.mean([masi(x, y) ** 2 for x, y in pairs]))
        allsets = [x for x, _ in pairs] + [y for _, y in pairs]
        De = float(np.mean([masi(allsets[i], allsets[j]) ** 2 for i in range(len(allsets)) for j in range(i + 1, len(allsets))])) if len(allsets) <= 1200 else float(np.mean([masi(allsets[i], allsets[j]) ** 2 for i, j in np.random.default_rng(44).integers(0, len(allsets), size=(200000, 2)) if i != j]))
        M[f"krippendorff_alpha_masi_{a}_vs_{b}"] = 1.0 - Do / De if De > 0 else None
        # Jaccard agreement on multi-label sets between raters
        M[f"mean_jaccard_{a}_vs_{b}"] = float(np.mean([len(set(A[w]["categories"]) & set(B[w]["categories"])) / max(len(set(A[w]["categories"]) | set(B[w]["categories"])), 1) for w in ids]))
        dis = [{"work_id": w, "title": df.set_index("work_id").title[w], "rules": sorted(rule[w]), "rule_primary": rule_p[w], a: sorted(A[w]["categories"]), b: sorted(B[w]["categories"]),
                "pi_review": bool(df.set_index("work_id").pi_review[w])} for w in ids if not (set(A[w]["categories"]) & set(B[w]["categories"]) & rule[w]) and (rule[w] or A[w]["categories"] or B[w]["categories"])]
        pd.DataFrame(dis).to_csv(f"{d}/disagreements.csv", index=False); M["n_disagreements"] = len(dis)
    json.dump(M, open(f"{d}/metrics.json", "w"), indent=1)
    print({k: v for k, v in M.items() if not isinstance(v, dict)}); print({k: (v["micro_f1"], v["macro_f1"]) for k, v in M.items() if isinstance(v, dict) and "micro_f1" in v})


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", default=None)
    ap.add_argument("--draw", action="store_true"); ap.add_argument("--rate", action="store_true"); ap.add_argument("--metrics", action="store_true")
    ap.add_argument("--model", default="claude-haiku-5-5"); ap.add_argument("--env-file", default="~/wp44_uqo/.env"); ap.add_argument("--seed", type=int, default=44)
    ap.add_argument("--per-cat", type=int, default=50); ap.add_argument("--n-none", type=int, default=200); ap.add_argument("--batch", type=int, default=10)
    a = ap.parse_args(); ont = a.ontology or os.path.join(a.root, "taxonomy/ontology.yaml")
    if a.draw: draw(a.root, a.seed, a.per_cat, a.n_none)
    if a.rate: rate(a.root, a.model, a.env_file, ont, a.batch)
    if a.metrics: metrics(a.root, ont)


if __name__ == "__main__":
    main()