"""F5-conf-02: Monte Carlo bias of OLS rho_hat in AR(1) with intercept.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py"""
import json, csv, os
import numpy as np
import 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"), dtype=str)
p = dict(zip(d["parameter"].str.strip(), d["value"].str.strip()))
n = int(p["n"]); rho = float(p["rho"]); sigma = float(p["sigma"])
R = int(p["n_replications"]); seed = int(p["seed"])
issues = []
assert p["intercept"].lower() == "yes" and p["initial_condition"].lower() == "stationary"
assert abs(rho) < 1

rng = np.random.default_rng(seed)
# Draw order: y_1 (stationary) for all reps, then innovations e_2..e_n (R x (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:]          # 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 ** 2).sum(axis=1)

# Cross-check a few replications against lstsq with explicit intercept
for i in range(5):
    X = np.column_stack([np.ones(n - 1), x[i]])
    b = np.linalg.lstsq(X, z[i], rcond=None)[0]
    assert abs(b[1] - rho_hat[i]) < 1e-10

finite = np.isfinite(rho_hat)
rh = rho_hat[finite]
Rk = int(rh.size)
with open(os.path.join(ROOT, "out", "exclusions.csv"), "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["step", "rule", "n_before", "n_after", "n_removed", "rationale"])
    w.writerow([1, "drop replications with non-finite rho_hat", R, Rk, R - Rk,
                "OLS undefined if regressor has zero variance (probability zero)"])

mean = float(rh.mean()); sd = float(rh.std(ddof=1))
bias_mc = mean - rho; se = sd / np.sqrt(Rk)
bias_an = -(1 + 3 * rho) / n
res = {
    "task_id": "F5-conf-02",
    "bias_analytical": bias_an,
    "bias_mc": bias_mc,
    "bias_mc_se": float(se),
    "rho_hat_mean": mean,
    "rho_hat_sd_mc": sd,
    "discrepancy_se_units": float((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": Rk,
    "random_seed": seed,
    # supplementary (not required) quantities
    "bias_analytical_neff": -(1 + 3 * rho) / (n - 1),
    "discrepancy_se_units_neff": float((bias_mc + (1 + 3 * rho) / (n - 1)) / se),
    "bias_mc_ci95_low": float(bias_mc - 1.96 * se),
    "bias_mc_ci95_high": float(bias_mc + 1.96 * se),
    "ratio_bias_mc_to_analytical": float(bias_mc / bias_an),
    "rho_hat_median": float(np.median(rh)),
    "share_rho_hat_below_rho": float(np.mean(rh < rho)),
    "se_asymptotic_over_sd_mc": float(np.sqrt((1 - rho**2) / n) / sd),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))