#!/usr/bin/env python3
"""WP44 — RQ2/H1/H3 driver: interrupted time series, placebo breaks, Bai–Perron, PELT and BOCPD on monthly series.

Series: total eligible works (month-precision records only), by area and by category (results/graph/category_counts_monthly.csv,
n_month_precision column when present; before classification only the total is run). Window: 2018-01 → --last-complete
(default 2026-06; sensitivity 2026-08 and 2026-03). Event dates: 2022-11-30 (ChatGPT, primary) and alternatives
2020-05 (GPT-3), 2022-04 (DALL·E 2), 2023-03 (GPT-4), 2024-09 (o1). Outputs results/econometrics/its_results.json,
placebo_<series>.csv, breaks.csv, and a per-series CSV of fitted values."""
import argparse, json, os, sys, numpy as np, pandas as pd
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from econometrics.structural_change import its_fit, placebo_breaks, bai_perron, pelt, bocpd

EVENTS = {"chatgpt": "2022-12-01", "gpt3": "2020-06-01", "dalle2": "2022-05-01", "gpt4": "2023-04-01", "o1": "2024-10-01"}


def series_total(R):
    d = pd.read_csv(f"{R}/results/bibliometrics/volume_monthly.csv")
    if "date_precision" in d.columns:
        d = d[d.date_precision == "month"]
    m = d.groupby("first_month")["n"].sum(); return {"total": m}


def series_categories(R):
    p = f"{R}/results/graph/category_counts_monthly.csv"
    if not os.path.exists(p):
        return {}
    d = pd.read_csv(p); col = "n_month_precision" if "n_month_precision" in d.columns else "n_multilabel"
    return {f"cat_{c}": g.groupby("first_month")[col].sum() for c, g in d.groupby("category_id")}


def run_one(name, s, first, last, out, placebo_baseline_only=True, donut=0):
    s = s[(s.index >= first) & (s.index <= last)].astype(float)
    idx = pd.period_range(first, last, freq="M").strftime("%Y-%m"); s = s.reindex(idx).fillna(0.0)
    if (s > 0).sum() < 48 or s.sum() < 500:
        return None
    y = s.values + 1.0; dates = pd.to_datetime(idx + "-01"); res = {"series": name, "n_months": len(y), "total": float(s.sum())}
    for ev, d in EVENTS.items():
        if pd.Timestamp(d) <= dates[12] or pd.Timestamp(d) >= dates[-12]:
            continue
        o, fit = its_fit(y, dates, d); res[f"its_{ev}"] = o
        if ev == "chatgpt":
            pl, df = placebo_breaks(y, dates, d); res["placebo_chatgpt"] = pl; df.to_csv(f"{out}/placebo_{name}.csv", index=False)
            if placebo_baseline_only:
                dfb = df[df.event < pd.Timestamp(d)]; tv = df[df.event == pd.Timestamp(d)]
                if len(tv) and len(dfb):
                    res["placebo_chatgpt_baseline_only"] = {"n_placebo": int(len(dfb)), **{f"{k}_placebo_p": float((np.abs(dfb[k]) >= abs(float(tv[k].iloc[0]))).mean()) for k in ("level_t", "slope_t")}}
            if donut:
                keep = np.abs((dates - pd.Timestamp(d)).days) > donut * 31
                od, _ = its_fit(y[keep], dates[keep], d); res["its_chatgpt_donut"] = {"months_excluded_each_side": donut, **od}
            pd.DataFrame({"month": idx, "y": s.values, "fitted": np.exp(fit.fittedvalues) - 1}).to_csv(f"{out}/fit_{name}.csv", index=False)
    ly = np.log(y); months_of_year = dates.month; base = dates < pd.Timestamp("2022-12-01")
    seas = {m: ly[base & (months_of_year == m)].mean() - ly[base].mean() for m in range(1, 13)}
    ly_sa = ly - np.array([seas[m] for m in months_of_year])   # seasonal adjustment from baseline month effects
    bp = bai_perron(ly_sa, dates, max_breaks=3); res["bai_perron"] = {k: v for k, v in bp.items() if k != "by_m"}; res["bai_perron"]["input"] = "log, seasonally adjusted (baseline month means)"
    res["bai_perron"]["bic_by_m"] = {m: round(v["bic"], 2) for m, v in bp["by_m"].items()}
    res["pelt"] = [idx[i] for i in pelt(ly_sa) if i < len(idx)]
    _, cp, mrl = bocpd(ly_sa); resets = [idx[t] for t in range(12, len(idx)) if mrl[t] <= 2 and mrl[t - 1] > 12]
    res["bocpd_resets"] = resets; res["bocpd_cp_prob_2022_12_to_2023_03"] = float(cp[[i for i, m in enumerate(idx) if "2022-12" <= m <= "2023-03"]].max()) if any("2022-12" <= m <= "2023-03" for m in idx) else None
    return res


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--first", default="2018-02"); ap.add_argument("--last-complete", default="2026-06")
    ap.add_argument("--sensitivity", default="2026-03,2026-08"); ap.add_argument("--donut", type=int, default=3); a = ap.parse_args(); R = a.root
    out = f"{R}/results/econometrics"; os.makedirs(out, exist_ok=True)
    S = series_total(R); S.update(series_categories(R)); all_res = {"last_complete": a.last_complete, "first": a.first, "events": EVENTS, "series": {}}
    for name, s in S.items():
        r = run_one(name, s, a.first, a.last_complete, out, donut=a.donut)
        if r:
            all_res["series"][name] = r; print(name, "ITS chatgpt level", round(r.get("its_chatgpt", {}).get("level_change", float("nan")), 3), "p", round(r.get("its_chatgpt", {}).get("level_p", float("nan")), 3), "placebo p", r.get("placebo_chatgpt", {}).get("level_t_placebo_p"), "BP", r["bai_perron"].get("break_dates_bic"), flush=True)
    sens = {}
    for lc in a.sensitivity.split(","):
        r = run_one("total", S["total"], a.first, lc, out + "/sens_tmp" if os.makedirs(out + "/sens_tmp", exist_ok=True) is None else out)
        if r:
            sens[lc] = {"its_chatgpt": r.get("its_chatgpt"), "bai_perron": r["bai_perron"].get("break_dates_bic")}
    all_res["sensitivity_last_complete_total"] = sens
    json.dump(all_res, open(f"{out}/its_results.json", "w"), indent=1, default=float)
    rows = [{"series": n, **{f"{ev}_level": r.get(f"its_{ev}", {}).get("level_change") for ev in EVENTS}, **{f"{ev}_level_p": r.get(f"its_{ev}", {}).get("level_p") for ev in EVENTS},
             "chatgpt_slope": r.get("its_chatgpt", {}).get("slope_change"), "chatgpt_slope_p": r.get("its_chatgpt", {}).get("slope_p"), "placebo_level_p": r.get("placebo_chatgpt", {}).get("level_t_placebo_p"),
             "placebo_slope_p": r.get("placebo_chatgpt", {}).get("slope_t_placebo_p"), "bp_breaks_bic": ";".join(r["bai_perron"].get("break_dates_bic", [])), "pelt": ";".join(r["pelt"]), "bocpd_resets": ";".join(r["bocpd_resets"])} for n, r in all_res["series"].items()]
    pd.DataFrame(rows).to_csv(f"{out}/breaks.csv", index=False); print("series analysed:", len(rows))


if __name__ == "__main__":
    main()