#!/usr/bin/env python3
"""WP44 — Phase 3: schema-constrained LLM extraction of method / model / dataset / benchmark / claim nodes and typed edges
(introduces_method, uses_dataset, evaluates_on, extends, contradicts, reproduces, addresses_problem) for the high-priority set,
from title + abstract (evidence_type = abstract) — full text is NOT used unless a lawful OA text is attached later.
Every edge carries provenance (work_id + quoted span from the abstract), confidence (model-reported, 0–1) and
verification_status = 'unverified' (machine-verified when the quoted span is found verbatim in the abstract; human-verified
only after PI review). Model: claude-sonnet-5-5 (structured outputs, effort low), batches of 5 works, resumable JSONL.
Budget: see protocol/amendments.csv (A-004)."""
import argparse, json, os, time, pandas as pd

SCHEMA = {"type": "object", "properties": {"works": {"type": "array", "items": {"type": "object", "properties": {
    "work_id": {"type": "integer"},
    "contribution_type": {"type": "string", "enum": ["theory", "method", "empirical", "benchmark", "engineering", "survey", "other"]},
    "entities": {"type": "array", "items": {"type": "object", "properties": {
        "name": {"type": "string"}, "kind": {"type": "string", "enum": ["method", "model", "dataset", "benchmark", "problem", "metric"]},
        "role": {"type": "string", "enum": ["introduced", "used", "evaluated_on", "extended", "compared_against", "addressed"]},
        "span": {"type": "string"}, "confidence": {"type": "number"}}, "required": ["name", "kind", "role", "span", "confidence"], "additionalProperties": False}},
    "claims": {"type": "array", "items": {"type": "object", "properties": {"claim": {"type": "string"}, "claim_kind": {"type": "string", "enum": ["result", "improvement", "negative_result", "reproduction", "contradiction"]},
        "span": {"type": "string"}, "confidence": {"type": "number"}}, "required": ["claim", "claim_kind", "span", "confidence"], "additionalProperties": False}}},
    "required": ["work_id", "contribution_type", "entities", "claims"], "additionalProperties": False}}}, "required": ["works"], "additionalProperties": False}
SYSTEM = ("You extract structured knowledge from computer-science paper abstracts for a knowledge graph. For each work: the contribution type; the named "
          "methods, models, datasets, benchmarks, metrics and research problems with their role; and the main claims. Every entity and claim must quote a "
          "short verbatim span (<= 25 words) from the abstract that supports it. Use canonical names (e.g. 'ImageNet', 'BERT', 'Transformer'). Never infer "
          "anything not stated in the abstract. Confidence is your calibrated probability that the extraction is correct.")


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--model", default="claude-sonnet-5-5"); ap.add_argument("--env-file", default="~/wp44_uqo/.env")
    ap.add_argument("--batch", type=int, default=5); ap.add_argument("--limit", type=int, default=0); ap.add_argument("--max-usd", type=float, default=150.0); a = ap.parse_args(); R = a.root
    import anthropic, duckdb
    key = next(l.split("=", 1)[1].strip() for l in open(os.path.expanduser(a.env_file)) if l.startswith("ANTHROPIC_API_KEY="))
    client = anthropic.Anthropic(api_key=key, max_retries=5)
    con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True)
    hp = pd.read_parquet(f"{R}/literature_synthesis/landmarks/high_priority.parquet")
    ab = con.execute(f"SELECT work_id, abstract FROM oa_abstracts WHERE work_id IN (SELECT work_id FROM read_parquet('{R}/literature_synthesis/landmarks/high_priority.parquet'))").df()
    hp = hp.merge(ab, on="work_id"); hp = hp.sort_values(["score"], ascending=False)
    if a.limit: hp = hp.head(a.limit)
    out_dir = f"{R}/data/processed/graph/extraction"; os.makedirs(out_dir, exist_ok=True); out_p = f"{out_dir}/extractions_{a.model}.jsonl"
    done = {json.loads(l)["work_id"] for l in open(out_p)} if os.path.exists(out_p) else set()
    todo = hp[~hp.work_id.isin(done)]; cost = 0.0; price_in, price_out, price_cache = 2.0 / 1e6, 10.0 / 1e6, 0.2 / 1e6; n_calls = 0
    with open(out_p, "a") as fh:
        for i in range(0, len(todo), a.batch):
            if cost >= a.max_usd:
                print("budget reached", cost); break
            b = todo.iloc[i:i + a.batch]
            items = "\n\n".join(f"### work_id {r.work_id}\nTitle: {r.title}\nYear: {r.year}\nAbstract: {(r.abstract or '')[:3000]}" for r in b.itertuples())
            for attempt in range(4):
                try:
                    resp = client.messages.create(model=a.model, max_tokens=12000, system=[{"type": "text", "text": SYSTEM, "cache_control": {"type": "ephemeral"}}],
                                                  messages=[{"role": "user", "content": items}], output_config={"format": {"type": "json_schema", "schema": SCHEMA}, "effort": "low"}); break
                except anthropic.RateLimitError:
                    time.sleep(15 * (attempt + 1))
            if resp.stop_reason == "refusal":
                continue
            u = resp.usage; cost += u.input_tokens * price_in + u.output_tokens * price_out + (u.cache_read_input_tokens or 0) * price_cache; n_calls += 1
            data = json.loads(next(bk.text for bk in resp.content if bk.type == "text"))
            abs_by_id = {int(r.work_id): (r.abstract or "") for r in b.itertuples()}
            for w in data["works"]:
                txt = abs_by_id.get(w["work_id"], "")
                for e in w["entities"] + w["claims"]:
                    e["verification_status"] = "machine-verified" if e["span"] and e["span"].strip().lower() in txt.lower() else "unverified"
                w["model"] = a.model; w["evidence_type"] = "abstract"; fh.write(json.dumps(w) + "\n")
            fh.flush()
            if n_calls % 25 == 0:
                print(f"{n_calls} calls, {i + len(b)}/{len(todo)} works, ${cost:.2f}", flush=True)
    json.dump({"model": a.model, "calls": n_calls, "cost_usd": round(cost, 2), "works_done": len(done) + min(len(todo), i + len(b) if len(todo) else 0)}, open(f"{out_dir}/extraction_usage_{a.model}.json", "w"), indent=1)
    print("done", n_calls, "calls", round(cost, 2), "USD")


if __name__ == "__main__":
    main()