"""F5-conf-02: finite-sample bias of OLS AR(1) coefficient (with intercept).
Reads only data/design.csv; writes out/results.json deterministically.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, math, 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"]))
data_issues = []

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

# Baseline / analytical quantities
bias_analytical = -(1 + 3 * rho) / n
se_asymptotic = math.sqrt((1 - rho ** 2) / n)

# Monte Carlo (vectorised across replications); c = 0 in DGP
rng = np.random.default_rng(seed)
y = np.empty((R, n))
y[:, 0] = rng.normal(0.0, sigma / math.sqrt(1 - rho ** 2), size=R)
eps = rng.normal(0.0, sigma, size=(R, n - 1))
for t in range(1, n):
    y[:, t] = rho * y[:, t - 1] + eps[:, t - 1]

x = y[:, :-1]; z = y[:, 1:]          # regress y_t on (1, y_{t-1}), t = 2..n
xd = x - x.mean(axis=1, keepdims=True)
zd = z - z.mean(axis=1, keepdims=True)
rho_hat = (xd * zd).sum(axis=1) / (xd * xd).sum(axis=1)

finite = np.isfinite(rho_hat)
if not finite.all():
    data_issues.append({"issue": f"{(~finite).sum()} non-finite rho_hat", "handling": "excluded"})
rh = rho_hat[finite]
R_used = int(rh.size)

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 / math.sqrt(R_used)
discrepancy = (bias_mc - bias_analytical) / bias_mc_se

# Sensitivity (not a required key): formula with effective sample size n-1
bias_analytical_nm1 = -(1 + 3 * rho) / (n - 1)

results = {
    "task_id": "F5-conf-02",
    "bias_analytical": bias_analytical,
    "bias_mc": bias_mc,
    "bias_mc_se": bias_mc_se,
    "rho_hat_mean": rho_hat_mean,
    "rho_hat_sd_mc": rho_hat_sd,
    "discrepancy_se_units": discrepancy,
    "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
    "baseline_naive_bias": 0.0,
    "bias_analytical_effective_n_minus_1": bias_analytical_nm1,
    "discrepancy_se_units_n_minus_1": (bias_mc - bias_analytical_nm1) / bias_mc_se,
    "rho_hat_median": float(np.median(rh)),
    "rho_hat_sd_ratio_to_asymptotic": rho_hat_sd / se_asymptotic,
    "relative_error_analytical_vs_mc": (bias_analytical - bias_mc) / bias_mc,
    "design": {"n": n, "rho": rho, "sigma": sigma, "intercept": True,
               "initial_condition": "stationary", "c": 0.0},
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(results, f, indent=2)
for k in ["bias_analytical", "bias_mc", "bias_mc_se", "discrepancy_se_units",
          "bias_analytical_effective_n_minus_1", "discrepancy_se_units_n_minus_1",
          "rho_hat_mean", "rho_hat_median", "rho_hat_sd_mc", "se_asymptotic",
          "relative_error_analytical_vs_mc"]:
    print(f"{k}: {results[k]:.6f}")