#!/usr/bin/env python3
"""WP44 — preregistered known-emergence set (protocol P1, H8) as keyword-defined monthly series (OpenAlex topics are too
coarse for these concepts: e.g. no topic for LLM agents, VLMs, speculative decoding, MoE). For each concept, monthly counts of
eligible month-precision works whose title or abstract matches the regex; detection = first month at which the trailing
3-month z-score against the trailing 24 months exceeds 3 AND 12-month growth > 100 % with >= 30 works; lead time =
peak-growth month (max 12-month log growth) minus detection month. Output: results/emerging/known_set_lead_times.csv."""
import argparse, os, re, duckdb, numpy as np, pandas as pd

KNOWN = {
    "large language models": r"\b(large language models?|LLMs?)\b",
    "diffusion models": r"\b(diffusion models?|denoising diffusion|latent diffusion|score-based generative)\b",
    "neural radiance fields": r"\b(neural radiance fields?|NeRFs?)\b",
    "llm agents": r"\b(LLM[- ]based agents?|language (model|agent)s? agents?|LLM agents?|agentic (AI|systems?|workflows?)|tool-using (agents?|LLMs?))\b",
    "retrieval-augmented generation": r"\b(retrieval[- ]augmented generation|\bRAG\b)",
    "state-space sequence models": r"\b(state[- ]space models?|Mamba|S4\b|selective state space|structured state space)\b",
    "vision-language models": r"\b(vision[- ]language models?|VLMs?|multimodal large language models?|MLLMs?)\b",
    "mechanistic interpretability": r"\b(mechanistic interpretability|sparse autoencoders? (for|in) (interpretab|LLM|language)|circuits? in (language|transformer))",
    "speculative decoding": r"\b(speculative (decoding|sampling)|draft models?)\b",
    "mixture of experts": r"\b(mixture[- ]of[- ]experts|\bMoE\b|sparsely[- ]gated)\b",
    "reasoning models": r"\b(chain[- ]of[- ]thought|test[- ]time (compute|scaling)|reasoning models?|large reasoning models?|o1-like|R1-like)\b",
    "ai code generation": r"\b(code generation|code LLMs?|program synthesis with (LLMs?|language models)|AI (pair|coding) (programmer|assistants?)|Copilot)\b",
}


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--last-month", default="2026-06"); a = ap.parse_args(); R = a.root
    out = f"{R}/results/emerging"; os.makedirs(out, exist_ok=True)
    con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True); con.execute("SET threads TO 16;")
    rows = []; series = {}
    for name, rx in KNOWN.items():
        q = f"""SELECT c.first_month AS m, count(*) n FROM corpus_works c LEFT JOIN oa_abstracts a ON a.work_id = c.id
                WHERE c.eligible AND c.date_precision = 'month' AND c.first_month BETWEEN '2018-01' AND '{a.last_month}'
                  AND (regexp_matches(c.title, ?) OR regexp_matches(coalesce(a.abstract, ''), ?)) GROUP BY 1 ORDER BY 1"""
        d = con.execute(q, [rx, rx]).df().set_index("m").n
        months = pd.period_range("2018-01", a.last_month, freq="M").strftime("%Y-%m"); s = d.reindex(months).fillna(0).astype(float); series[name] = s
        v = s.values; det = None; g12 = np.full(len(v), np.nan)
        for t in range(24, len(v)):
            last3 = v[t - 2:t + 1].mean(); trail = v[t - 26:t - 2]; z = (last3 - trail.mean()) / (trail.std() + 1e-6)
            v12 = v[t - 11:t + 1].sum(); v12p = v[t - 23:t - 11].sum(); g12[t] = np.log1p(v12) - np.log1p(v12p)
            if det is None and z > 3 and v12 >= 30 and v12 >= 2 * v12p:
                det = months[t]
        peak = months[int(np.nanargmax(g12))] if np.isfinite(g12).any() else None
        lead = (pd.Period(peak, "M") - pd.Period(det, "M")).n if det and peak else None
        rows.append({"concept": name, "regex": rx, "total_works": int(v.sum()), "first_month_with_5": next((m for m, x in zip(months, np.cumsum(v)) if x >= 5), None),
                     "detection_month": det, "peak_growth_month": peak, "lead_months": lead})
    pd.DataFrame(rows).to_csv(f"{out}/known_set_lead_times.csv", index=False); pd.DataFrame(series).to_csv(f"{out}/known_set_monthly.csv")
    print(pd.DataFrame(rows)[["concept", "total_works", "detection_month", "peak_growth_month", "lead_months"]])


if __name__ == "__main__":
    main()