"""F5-conf-02: finite-sample bias of OLS AR(1) coefficient (with intercept).
Reads data/design.csv only; writes out/results.json deterministically.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np
import pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
design = pd.read_csv(os.path.join(ROOT, "data", "design.csv"))
p = dict(zip(design["parameter"], design["value"]))

n = int(p["n"]); rho = float(p["rho"]); sigma = float(p["sigma"])
intercept = str(p["intercept"]).strip().lower() == "yes"
init = str(p["initial_condition"]).strip().lower()
R = int(p["n_replications"]); seed = int(p["seed"])
data_issues = []
assert intercept and init == "stationary" and abs(rho) < 1 and n >= 3 and R > 1

# --- Baseline (closed-form) quantities ---
bias_analytical = -(1.0 + 3.0 * rho) / n
se_asymptotic = float(np.sqrt((1.0 - rho ** 2) / n))
# sensitivity: same formula using effective regression sample size n-1
bias_analytical_neff = -(1.0 + 3.0 * rho) / (n - 1)

# --- Monte Carlo (vectorised) ---
# Draw order: first R stationary initial values, then an R x (n-1) innovation matrix.
rng = np.random.default_rng(seed)
y = np.empty((R, n))
y[:, 0] = rng.standard_normal(R) * sigma / np.sqrt(1.0 - rho ** 2)
eps = rng.standard_normal((R, n - 1)) * sigma
c = 0.0
for t in range(1, n):
    y[:, t] = c + rho * y[:, t - 1] + eps[:, t - 1]
x = y[:, :-1]; z = y[:, 1:]           # regress y_t (t=2..n) on (1, y_{t-1})
xd = x - x.mean(axis=1, keepdims=True)
zd = z - z.mean(axis=1, keepdims=True)
rho_hat = (xd * zd).sum(axis=1) / (xd ** 2).sum(axis=1)
finite = np.isfinite(rho_hat)
n_nonfinite = int((~finite).sum())
if n_nonfinite:
    data_issues.append({"issue": f"{n_nonfinite} non-finite rho_hat", "handling": "excluded"})
rh = rho_hat[finite]
R_used = int(rh.size)

# cross-check one replication against numpy lstsq
X0 = np.column_stack([np.ones(n - 1), x[0]])
b0 = np.linalg.lstsq(X0, z[0], rcond=None)[0][1]
assert abs(b0 - rho_hat[0]) < 1e-10

rho_hat_mean = float(rh.mean())
rho_hat_sd = float(rh.std(ddof=1))
bias_mc = rho_hat_mean - rho
bias_mc_se = rho_hat_sd / np.sqrt(R_used)
disc = (bias_mc - bias_analytical) / bias_mc_se

res = {
    "task_id": "F5-conf-02",
    "bias_analytical": bias_analytical,
    "bias_mc": float(bias_mc),
    "bias_mc_se": float(bias_mc_se),
    "rho_hat_mean": rho_hat_mean,
    "rho_hat_sd_mc": rho_hat_sd,
    "discrepancy_se_units": float(disc),
    "se_asymptotic": se_asymptotic,
    "n_replications": R,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"design.csv": int(len(design))},
    "n_analysed": R_used,
    "random_seed": seed,
    # supplementary keys
    "baseline_bias_analytical": bias_analytical,
    "baseline_se_asymptotic": se_asymptotic,
    "bias_analytical_neff": bias_analytical_neff,
    "discrepancy_se_units_neff": float((bias_mc - bias_analytical_neff) / bias_mc_se),
    "bias_mc_minus_analytical": float(bias_mc - bias_analytical),
    "ratio_bias_mc_to_analytical": float(bias_mc / bias_analytical),
    "rho_hat_median": float(np.median(rh)),
    "share_rho_hat_below_rho": float((rh < rho).mean()),
    "ratio_sd_mc_to_se_asymptotic": float(rho_hat_sd / se_asymptotic),
    "n_regression_obs": n - 1,
    "n_nonfinite_rho_hat": n_nonfinite,
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))