"""F2-conf-02: OLS slope of y on x with Newey-West HAC SE. Run from workspace root."""
import json, math
import numpy as np, pandas as pd, statsmodels.api as sm
from scipy import stats

d = pd.read_csv("data/series.csv")
n_rows = len(d)
issues = []
# data checks
months = pd.to_datetime(d["month"], format="%Y-%m")
assert d["month"].is_unique and d[["x", "y"]].notna().all().all()
gaps = (months.dt.year * 12 + months.dt.month).diff().dropna()
assert (gaps == 1).all(), "non-consecutive months"
d = d.sort_values("month").reset_index(drop=True)

X = sm.add_constant(d["x"])
ols = sm.OLS(d["y"], X).fit()
alpha_hat, beta_hat = float(ols.params["const"]), float(ols.params["x"])
se_iid = float(ols.bse["x"])
e = ols.resid.values
rho_resid = float(np.sum(e[1:] * e[:-1]) / np.sum(e ** 2))  # lag-1 autocorrelation (mean-zero residuals)

# Bandwidth: Andrews (1991) AR(1) plug-in for the Bartlett kernel, applied to the
# slope score v_t = (x_t - xbar) * u_t, which is what the long-run variance concerns.
T = len(d)
v = (d["x"].values - d["x"].values.mean()) * e
rho_v = float(np.sum(v[1:] * v[:-1]) / np.sum(v[:-1] ** 2))
a1 = 4 * rho_v ** 2 / ((1 - rho_v) ** 2 * (1 + rho_v) ** 2)
hac_lags = int(math.floor(1.1447 * (a1 * T) ** (1 / 3)))
nw_rule = int(math.floor(4 * (T / 100) ** (2 / 9)))  # Newey-West (1994) rule of thumb, for sensitivity

def hac(L):
    r = ols.get_robustcov_results(cov_type="HAC", maxlags=L, use_correction=True)
    return float(r.bse[1])

se_hac = hac(hac_lags)
t_hac = beta_hat / se_hac
p_hac = float(2 * stats.norm.sf(abs(t_hac)))  # asymptotic normal reference
z = stats.norm.ppf(0.975)
ci = (beta_hat - z * se_hac, beta_hat + z * se_hac)

sens = {}
for L in sorted(set([0, nw_rule, hac_lags, 12, 24, 36])):
    s = hac(L)
    sens[str(L)] = {"se": s, "t": beta_hat / s, "p_normal": float(2 * stats.norm.sf(abs(beta_hat / s))),
                    "p_t": float(2 * stats.t.sf(abs(beta_hat / s), T - 2))}

res = {
    "task_id": "F2-conf-02",
    "beta_hat": beta_hat, "alpha_hat": alpha_hat, "se_iid": se_iid,
    "t_iid": beta_hat / se_iid, "p_iid": float(ols.pvalues["x"]),
    "rho_resid": rho_resid, "rho_score": rho_v,
    "se_hac": se_hac, "hac_lags": hac_lags, "t_hac": t_hac, "p_hac": p_hac,
    "p_hac_t_dist": float(2 * stats.t.sf(abs(t_hac), T - 2)),
    "beta_ci_low": ci[0], "beta_ci_high": ci[1],
    "memo_claim_supported": bool(p_hac < 0.05),
    "n_analysed": int(T),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"series.csv": n_rows},
    "random_seed": None,
    "hac_sensitivity_by_lags": sens, "nw_rule_of_thumb_lags": nw_rule,
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))