#!/usr/bin/env python3
"""WP44 — Phase 6: emerging-research detection framework (RQ10, hypothesis H8), strictly out-of-time.

Units: OpenAlex topics (T-ids) restricted to those with >= 50 corpus works over 2018–2026 (topic-level series are the
most reproducible unit; semantic clusters are an extension in experiments/emerging_clusters.py).
Monthly panel from the eligible corpus (month-precision records only): n_{k,t} works per topic k and month t; citation
inflow c_{k,t} (internal citations received in month t by works of topic k, from cit_edges_internal with citing first
month); cross-area share x_{k,t}; mean novelty (optional, from embeddings) z_{k,t}.

Signals at cutoff T (use data <= T only; all means/sd computed on the trailing window ending at T):
  g12  growth of 12-month volume vs previous 12 months (log ratio)
  acc  acceleration: g12 at T minus g12 at T-12
  zvol z-score of the last 3 months' volume against the trailing 24-month mean/sd
  cit_acc  log ratio of citation inflow last 12 months vs previous 12
  xshare   cross-area share of references in the last 12 months (interdisciplinary connectivity)
  share    topic share of corpus volume (size control)
  age      months since first appearance (>= 5 works)
Label at horizon h (12 or 24 months): emerging_{k,T,h} = 1 if the topic's growth g_h over (T, T+h] is in the top decile
among topics with >= 20 works in (T-12, T], AND its absolute volume in (T, T+h] >= 2x the volume in (T-h, T].
Models: z-score baseline (rank by zvol), logistic regression (L2) and gradient boosting trained on cutoffs
<= 2021-12 (labels observed by 2023-12) and tested on cutoffs 2023-01 … 2024-10 (labels observed by 2026-10).
Metrics: AUC, precision@20/50, lead time on the preregistered known-emergence set (protocol P1, H8).
Outputs: results/emerging/{panel.parquet, features.parquet, metrics.json, predictions.csv, known_set_lead_times.csv}.
"""
import argparse, json, os, duckdb, numpy as np, pandas as pd

KNOWN_SET = {  # preregistered (protocol P1, H8): OpenAlex topic ids mapped after taxonomy work; filled by experiments/known_set.csv
}


def build_panel(root, con, last_month):
    con.execute(f"""CREATE OR REPLACE TABLE em_panel AS
        WITH W AS (SELECT c.id, c.first_month AS m, t.topic_id FROM corpus_works c JOIN oa_work_topics t ON t.work_id = c.id AND t.rank = 0
                   WHERE c.eligible AND c.date_precision = 'month' AND c.first_month BETWEEN '2018-01' AND '{last_month}'),
             V AS (SELECT topic_id, m, count(*) n FROM W GROUP BY 1,2),
             C AS (SELECT w.topic_id, cw.first_month AS m, count(*) c FROM cit_edges_internal e JOIN W w ON w.id = e.cited_id
                   JOIN corpus_works cw ON cw.id = e.citing_id WHERE cw.date_precision = 'month' GROUP BY 1,2)
        SELECT v.topic_id, v.m, v.n, coalesce(c.c, 0) AS c FROM V v LEFT JOIN C c USING (topic_id, m)""")
    df = con.execute("SELECT * FROM em_panel").df()
    keep = df.groupby("topic_id")["n"].sum(); keep = keep[keep >= 50].index
    df = df[df.topic_id.isin(keep)]
    months = pd.period_range("2018-01", last_month, freq="M").strftime("%Y-%m")
    P = df.pivot_table(index="topic_id", columns="m", values="n", fill_value=0).reindex(columns=months, fill_value=0)
    C = df.pivot_table(index="topic_id", columns="m", values="c", fill_value=0).reindex(columns=months, fill_value=0)
    return P, C


def features_at(P, C, T_idx):
    """Signals computed with columns <= T_idx only."""
    n = P.values; c = C.values
    def win(a, lo, hi):  # sum over columns (lo, hi]
        lo = max(lo, -1); return a[:, lo + 1:hi + 1].sum(1)
    v12 = win(n, T_idx - 12, T_idx); v12p = win(n, T_idx - 24, T_idx - 12); v12pp = win(n, T_idx - 36, T_idx - 24)
    g12 = np.log1p(v12) - np.log1p(v12p); g12p = np.log1p(v12p) - np.log1p(v12pp)
    last3 = win(n, T_idx - 3, T_idx) / 3.0
    trail = n[:, max(0, T_idx - 27):T_idx - 3 + 1]
    zvol = (last3 - trail.mean(1)) / (trail.std(1) + 1e-6)
    c12 = win(c, T_idx - 12, T_idx); c12p = win(c, T_idx - 24, T_idx - 12)
    cit_acc = np.log1p(c12) - np.log1p(c12p)
    cum = n[:, :T_idx + 1].cumsum(1); first = (cum >= 5).argmax(1); age = np.where(cum[:, -1] >= 5, T_idx - first, -1)
    share = v12 / max(v12.sum(), 1)
    return pd.DataFrame({"topic_id": P.index, "T_idx": T_idx, "v12": v12, "g12": g12, "acc": g12 - g12p, "zvol": zvol, "cit_acc": cit_acc,
                         "share": share, "age": age, "c12": c12})


def labels_at(P, T_idx, h):
    n = P.values
    pre = n[:, max(0, T_idx - h + 1):T_idx + 1].sum(1); post = n[:, T_idx + 1:T_idx + h + 1].sum(1)
    base12 = n[:, max(0, T_idx - 11):T_idx + 1].sum(1)
    g = np.log1p(post) - np.log1p(pre)
    elig = base12 >= 20
    thr = np.quantile(g[elig], 0.9) if elig.sum() > 10 else np.inf
    lab = ((g >= thr) & (post >= 2 * np.maximum(pre, 1)) & elig).astype(int)
    return pd.DataFrame({"topic_id": P.index, "T_idx": T_idx, "label": lab, "g_future": g, "eligible": elig})


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--last-month", default="2026-10")
    ap.add_argument("--horizon", type=int, default=12); a = ap.parse_args(); R = a.root
    out = os.path.join(R, "results/emerging"); os.makedirs(out, exist_ok=True)
    con = duckdb.connect(os.path.join(R, "data/processed/wp44.duckdb")); con.execute("SET threads TO 20; SET memory_limit='60GB';")
    P, C = build_panel(R, con, a.last_month); P.to_parquet(f"{out}/panel_volume.parquet"); C.to_parquet(f"{out}/panel_citations.parquet")
    months = list(P.columns); h = a.horizon
    def idx(m): return months.index(m)
    last_train = "2022-09" if h <= 12 else "2021-09"   # labels of training cutoffs must be fully observed before the first test cutoff
    train_cuts = [idx(m) for m in months if "2019-12" <= m <= last_train and idx(m) % 3 == 0]
    test_cuts = [idx(m) for m in months if "2023-01" <= m <= "2024-10" and idx(m) % 3 == 0 and idx(m) + h < len(months)]
    rows = []
    for T in train_cuts + test_cuts:
        f = features_at(P, C, T); l = labels_at(P, T, h); rows.append(f.merge(l, on=["topic_id", "T_idx"]))
    D = pd.concat(rows); D["split"] = np.where(D.T_idx.isin(train_cuts), "train", "test"); D = D[D.eligible]
    D.to_parquet(f"{out}/features.parquet")
    feats = ["g12", "acc", "zvol", "cit_acc", "share", "age"]
    from sklearn.linear_model import LogisticRegression
    from sklearn.ensemble import GradientBoostingClassifier
    from sklearn.metrics import roc_auc_score
    tr = D[D.split == "train"]; te = D[D.split == "test"]
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler().fit(tr[feats].replace([np.inf, -np.inf], 0).fillna(0).values)   # fitted on training cutoffs only (no look-ahead)
    X = lambda d: scaler.transform(d[feats].replace([np.inf, -np.inf], 0).fillna(0).values)
    models = {"zscore_baseline": None, "logistic": LogisticRegression(C=1.0, max_iter=2000), "gbm": GradientBoostingClassifier(n_estimators=300, max_depth=3, learning_rate=0.05, random_state=44)}
    metrics = {"horizon": h, "n_train": int(len(tr)), "n_test": int(len(te)), "pos_rate_train": float(tr.label.mean()), "pos_rate_test": float(te.label.mean()),
               "train_cutoffs": [months[i] for i in train_cuts], "test_cutoffs": [months[i] for i in test_cuts]}
    preds = te[["topic_id", "T_idx", "label", "g_future"]].copy()
    for name, m in models.items():
        if m is None:
            s = te["zvol"].values
        else:
            m.fit(X(tr), tr.label); s = m.predict_proba(X(te))[:, 1]
        preds[name] = s
        res = {"auc": float(roc_auc_score(te.label, s))}
        for k in (20, 50):
            pk = []
            for T, g in preds.groupby("T_idx"):
                top = g.nlargest(k, name); pk.append(top.label.mean())
            res[f"precision_at_{k}"] = float(np.mean(pk))
        metrics[name] = res
    # per-cutoff AUC stability
    metrics["auc_by_cutoff"] = {months[T]: {n: float(roc_auc_score(g.label, g[n])) if g.label.nunique() > 1 else None for n in models} for T, g in preds.groupby("T_idx")}
    preds["month"] = [months[i] for i in preds.T_idx]; preds.to_csv(f"{out}/predictions.csv", index=False)
    json.dump(metrics, open(f"{out}/metrics_h{h}.json", "w"), indent=1); print(json.dumps({k: v for k, v in metrics.items() if k != "auc_by_cutoff"}, indent=1))


if __name__ == "__main__":
    main()