"""F5-conf-02: Monte Carlo finite-sample bias of OLS AR(1) estimator vs Marriott-Pope/Kendall approximation.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, os
import numpy as np, pandas as pd

d = pd.read_csv("data/design.csv")
n_rows = len(d)
p = dict(zip(d["parameter"].astype(str).str.strip(), d["value"].astype(str).str.strip()))
issues = []
n = int(p["n"]); rho = float(p["rho"]); sigma = float(p["sigma"])
R = int(p["n_replications"]); seed = int(p["seed"])
intercept = p["intercept"].lower() in ("yes", "true", "1")
init = p["initial_condition"].lower()
assert intercept and init == "stationary" and abs(rho) < 1 and sigma > 0 and n > 2

rng = np.random.default_rng(seed)
# Draw all shocks as one (R, n) standard-normal array; column 0 -> y_1 from the stationary law.
z = rng.standard_normal((R, n))
y = np.empty((R, n))
y[:, 0] = z[:, 0] * sigma / np.sqrt(1 - rho**2)
for t in range(1, n):
    y[:, t] = rho * y[:, t - 1] + sigma * z[:, t]   # c = 0
# OLS of y_t on (1, y_{t-1}), t = 2..n  (n-1 usable observations)
x = y[:, :-1]; yy = y[:, 1:]
xc = x - x.mean(axis=1, keepdims=True); yc = yy - yy.mean(axis=1, keepdims=True)
rho_hat = (xc * yc).sum(axis=1) / (xc**2).sum(axis=1)
assert np.all(np.isfinite(rho_hat))

mean = float(rho_hat.mean()); sd = float(rho_hat.std(ddof=1))
bias_mc = mean - rho; se = sd / np.sqrt(R)
bias_an = -(1 + 3 * rho) / n
res = {
    "task_id": "F5-conf-02",
    "bias_analytical": bias_an,
    "bias_mc": bias_mc,
    "bias_mc_se": se,
    "rho_hat_mean": mean,
    "rho_hat_sd_mc": sd,
    "discrepancy_se_units": (bias_mc - bias_an) / se,
    "se_asymptotic": float(np.sqrt((1 - rho**2) / n)),
    "n_replications": R,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"design.csv": n_rows},
    "n_analysed": int(rho_hat.size),
    "random_seed": seed,
    # supplementary (not required) diagnostics
    "supplementary": {
        "bias_analytical_using_n_minus_1": -(1 + 3 * rho) / (n - 1),
        "median_rho_hat": float(np.median(rho_hat)),
        "share_rho_hat_below_rho": float((rho_hat < rho).mean()),
        "bias_mc_ci95": [bias_mc - 1.96 * se, bias_mc + 1.96 * se],
        "ratio_bias_mc_to_analytical": bias_mc / bias_an,
        "numpy_version": np.__version__,
    },
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))