"""F5-conf-02: Monte Carlo finite-sample bias of OLS AR(1) estimator vs Marriott-Pope/Kendall approximation."""
import json, os
import numpy as np, pandas as pd

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
d = pd.read_csv(os.path.join(ROOT, "data", "design.csv"))
p = dict(zip(d["parameter"], d["value"]))
issues = []
n = int(float(p["n"])); rho = float(p["rho"]); sigma = float(p["sigma"])
R = int(float(p["n_replications"])); seed = int(float(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 > 1

rng = np.random.default_rng(seed)
# Draw order: y_1 stationary draws (R), then innovations (R, n-1)
y = np.empty((R, n))
y[:, 0] = rng.normal(0.0, sigma / np.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]   # c = 0

x = y[:, :-1]; z = y[:, 1:]                       # t = 2..n
xc = x - x.mean(axis=1, keepdims=True)
zc = z - z.mean(axis=1, keepdims=True)
rho_hat = (xc * zc).sum(1) / (xc * xc).sum(1)     # OLS slope with intercept
ok = np.isfinite(rho_hat)
if not ok.all():
    issues.append({"issue": f"{(~ok).sum()} non-finite estimates", "handling": "excluded"})
rh = rho_hat[ok]; Ru = rh.size

mean = rh.mean(); sd = rh.std(ddof=1)
bias_mc = mean - rho; se = sd / np.sqrt(Ru)
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": int(len(d))},
    "n_analysed": int(Ru),
    "random_seed": seed,
    # supplementary diagnostics
    "supplementary": {
        "bias_analytical_using_n_minus_1": -(1 + 3 * rho) / (n - 1),
        "rho_hat_median": float(np.median(rh)),
        "share_rho_hat_below_rho": float((rh < rho).mean()),
        "ratio_bias_mc_to_analytical": bias_mc / bias_an,
        "ratio_sd_mc_to_se_asymptotic": sd / np.sqrt((1 - rho**2) / n),
    },
}
res = {k: (float(v) if isinstance(v, np.floating) else v) for k, v in res.items()}
res["supplementary"] = {k: float(v) for k, v in res["supplementary"].items()}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))