#!/usr/bin/env python3
"""WP44 — structural-change toolkit for monthly publication series (RQ2, hypotheses H1/H3).

Functions (all pure, tested in tests/test_structural_change.py):
  its_fit(y, dates, event, lags)      interrupted time series on log counts with month-of-year seasonality,
                                      level and slope change at `event`, Newey–West HAC standard errors
  placebo_breaks(y, dates, window)    distribution of the ITS level/slope t-statistics over every candidate month of the
                                      baseline (placebo), returns the empirical rank/p-value of the true event
  bai_perron(y, dates, max_breaks)    multiple structural breaks in mean + trend by dynamic programming (BIC), with
                                      minimum segment length h (Bai & Perron 1998/2003, pure-python DP on SSR)
  pelt(y, pen)                        PELT changepoints (ruptures, l2 cost) on log counts, penalty by BIC
  bocpd(y, hazard, mu0, kappa0, ...)  Bayesian online changepoint detection (Adams & MacKay 2007), Student-t predictive,
                                      returns run-length posterior and P(changepoint at t)
Right-truncation: callers must stop the series at the last complete month (see bibliometrics/01_volume_series.py).
"""
from __future__ import annotations
import numpy as np, pandas as pd
import statsmodels.api as sm


def _design(dates: pd.DatetimeIndex, event: pd.Timestamp, seasonal=True):
    t = np.arange(len(dates), dtype=float)
    post = (dates >= event).astype(float)
    t_post = np.where(post == 1, t - t[post == 1][0] if post.sum() else 0, 0.0)
    X = {"const": np.ones(len(t)), "t": t, "post": post, "t_post": t_post}
    if seasonal:
        for m in range(2, 13):
            X[f"m{m}"] = (dates.month == m).astype(float)
    return pd.DataFrame(X, index=dates)


def its_fit(y, dates, event, lags=6, seasonal=True, log=True):
    dates = pd.DatetimeIndex(dates); event = pd.Timestamp(event)
    yy = np.log(np.asarray(y, dtype=float)) if log else np.asarray(y, dtype=float)
    X = _design(dates, event, seasonal)
    res = sm.OLS(yy, X).fit(cov_type="HAC", cov_kwds={"maxlags": lags})
    out = {"level_change": res.params["post"], "level_se": res.bse["post"], "level_t": res.tvalues["post"], "level_p": res.pvalues["post"],
           "slope_change": res.params["t_post"], "slope_se": res.bse["t_post"], "slope_t": res.tvalues["t_post"], "slope_p": res.pvalues["t_post"],
           "pre_slope": res.params["t"], "pre_slope_se": res.bse["t"], "n": int(res.nobs), "r2": res.rsquared, "event": str(event.date())}
    return out, res


def placebo_breaks(y, dates, true_event, min_pre=24, min_post=12, lags=6, seasonal=True):
    dates = pd.DatetimeIndex(dates)
    rows = []
    for i in range(min_pre, len(dates) - min_post):
        ev = dates[i]
        o, _ = its_fit(y, dates, ev, lags, seasonal)
        rows.append({"event": ev, "level_t": o["level_t"], "slope_t": o["slope_t"], "level_change": o["level_change"], "slope_change": o["slope_change"]})
    df = pd.DataFrame(rows)
    tv = df.loc[df.event == pd.Timestamp(true_event)]
    out = {"n_placebo": len(df)}
    if len(tv):
        for k in ("level_t", "slope_t"):
            v = float(tv[k].iloc[0]); out[f"{k}_true"] = v
            out[f"{k}_placebo_p"] = float((np.abs(df[k]) >= abs(v)).mean())
            out[f"{k}_rank"] = int((np.abs(df[k]) >= abs(v)).sum())
    return out, df


def bai_perron(y, dates=None, max_breaks=3, h=None, trend=True):
    """Global minimiser of SSR with m breaks (Bai & Perron 2003 dynamic programming) for a segment-wise linear trend
    (or mean) model; model selection by BIC and by the LWZ criterion. Returns dict with breaks (indices and dates)."""
    y = np.asarray(y, dtype=float); n = len(y); h = h or max(8, int(0.15 * n))
    t = np.arange(n, dtype=float)
    def ssr(i, j):  # segment [i, j)
        yy = y[i:j]
        if trend:
            X = np.column_stack([np.ones(j - i), t[i:j]])
            beta, *_ = np.linalg.lstsq(X, yy, rcond=None); r = yy - X @ beta
        else:
            r = yy - yy.mean()
        return float(r @ r)
    S = np.full((n + 1, n + 1), np.inf)
    for i in range(n):
        for j in range(i + h, n + 1):
            S[i, j] = ssr(i, j)
    best = {}
    k_params = 2 if trend else 1
    # DP: C[m][j] = min SSR of first j obs with m breaks
    C = {0: S[0, :].copy()}; B = {0: {j: [] for j in range(n + 1)}}
    for m in range(1, max_breaks + 1):
        C[m] = np.full(n + 1, np.inf); B[m] = {}
        for j in range((m + 1) * h, n + 1):
            cands = [(C[m - 1][i] + S[i, j], i) for i in range(m * h, j - h + 1) if np.isfinite(C[m - 1][i]) and np.isfinite(S[i, j])]
            if cands:
                v, i = min(cands); C[m][j] = v; B[m][j] = B[m - 1][i] + [i]
    for m in range(0, max_breaks + 1):
        tot = C[m][n]
        if not np.isfinite(tot):
            continue
        p = (m + 1) * k_params + m
        bic = n * np.log(tot / n) + p * np.log(n)
        lwz = n * np.log(tot / (n - p)) + p * 0.299 * np.log(n) ** 2.1
        best[m] = {"ssr": tot, "bic": bic, "lwz": lwz, "breaks": B[m].get(n, [])}
    m_bic = min(best, key=lambda m: best[m]["bic"]); m_lwz = min(best, key=lambda m: best[m]["lwz"])
    res = {"by_m": best, "m_bic": m_bic, "m_lwz": m_lwz, "breaks_bic": best[m_bic]["breaks"], "h": h}
    if dates is not None:
        d = pd.DatetimeIndex(dates); res["break_dates_bic"] = [str(d[i].date()) for i in best[m_bic]["breaks"]]
        res["break_dates_lwz"] = [str(d[i].date()) for i in best[m_lwz]["breaks"]]
    return res


def pelt(y, pen=None, min_size=6, model="l2"):
    import ruptures as rpt
    y = np.asarray(y, dtype=float); n = len(y)
    pen = pen if pen is not None else 2 * np.log(n) * np.var(y)  # BIC-type penalty for l2 cost
    algo = rpt.Pelt(model=model, min_size=min_size, jump=1).fit(y.reshape(-1, 1))
    return algo.predict(pen=pen)[:-1]


def bocpd(y, hazard=1 / 60.0, mu0=None, kappa0=1.0, alpha0=1.0, beta0=None):
    """Adams & MacKay (2007) with Normal–Inverse-Gamma conjugacy (Student-t predictive). Returns (R, cp_prob, map_rl):
    R[t, r] is the run-length posterior after observing y[:t]; with a constant hazard the mass on r = 0 equals the hazard
    by construction, so the detection statistics are cp_prob[t] = P(r_t <= 3 | y[:t+1]) (a changepoint in the last
    months) and map_rl[t] = argmax_r R[t+1, r] (the MAP run length, which resets after a change)."""
    from scipy.stats import t as student
    y = np.asarray(y, dtype=float); n = len(y)
    mu0 = y[:6].mean() if mu0 is None else mu0; beta0 = y[:6].var() + 1e-6 if beta0 is None else beta0
    R = np.zeros((n + 1, n + 1)); R[0, 0] = 1.0
    mu = np.array([mu0]); kappa = np.array([kappa0]); alpha = np.array([alpha0]); beta = np.array([beta0])
    cp = np.zeros(n); map_rl = np.zeros(n, dtype=int)
    for t_ in range(n):
        x = y[t_]
        df = 2 * alpha; scale = np.sqrt(beta * (kappa + 1) / (alpha * kappa))
        pred = student.pdf(x, df, loc=mu, scale=scale)
        growth = R[t_, : t_ + 1] * pred * (1 - hazard)
        change = (R[t_, : t_ + 1] * pred * hazard).sum()
        R[t_ + 1, 1 : t_ + 2] = growth; R[t_ + 1, 0] = change
        R[t_ + 1] /= R[t_ + 1].sum() + 1e-300
        cp[t_] = R[t_ + 1, :4].sum(); map_rl[t_] = int(np.argmax(R[t_ + 1, : t_ + 2]))
        mu_new = (kappa * mu + x) / (kappa + 1); kappa_new = kappa + 1; alpha_new = alpha + 0.5
        beta_new = beta + kappa * (x - mu) ** 2 / (2 * (kappa + 1))
        mu = np.concatenate([[mu0], mu_new]); kappa = np.concatenate([[kappa0], kappa_new]); alpha = np.concatenate([[alpha0], alpha_new]); beta = np.concatenate([[beta0], beta_new])
    return R, cp, map_rl