"""F2-smoke-01: OLS slope of y on x with firm-clustered SE. Run from workspace root."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

d_raw = pd.read_csv("data/panel.csv")
n_raw = len(d_raw)
d = d_raw.dropna(subset=["firm_id", "x", "y"])
issues = []
if len(d) < n_raw:
    issues.append({"issue": f"{n_raw-len(d)} rows with missing values", "handling": "dropped (complete-case)"})
if d.duplicated().sum() > 0:
    issues.append({"issue": "exact duplicate rows", "handling": "retained (no basis to treat as errors)"})

y = d["y"].to_numpy(float); x = d["x"].to_numpy(float); g = d["firm_id"].to_numpy()
n = len(y); X = np.column_stack([np.ones(n), x]); k = 2
XtX_inv = np.linalg.inv(X.T @ X)
b = XtX_inv @ X.T @ y
e = y - X @ b
# conventional homoskedastic SE (naive baseline)
s2 = e @ e / (n - k)
se_ols = float(np.sqrt(s2 * XtX_inv[1, 1]))
# cluster-robust (CR1, Stata-style small-sample factor), clustered by firm_id
groups = np.unique(g); G = len(groups)
meat = np.zeros((k, k))
for gg in groups:
    m = g == gg
    s = X[m].T @ e[m]
    meat += np.outer(s, s)
c = G / (G - 1) * (n - 1) / (n - k)
V = c * XtX_inv @ meat @ XtX_inv
se_cl = float(np.sqrt(V[1, 1]))
tcrit = float(stats.t.ppf(0.975, G - 1))
res = {
    "task_id": "F2-smoke-01",
    "beta_hat": float(b[1]), "alpha_hat": float(b[0]),
    "se_ols": se_ols, "se_cluster": se_cl,
    "beta_ci_low": float(b[1] - tcrit * se_cl), "beta_ci_high": float(b[1] + tcrit * se_cl),
    "n_analysed": int(n), "n_clusters": int(G),
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": issues, "n_input_rows": {"panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary
    "ci_t_critical_df_Gminus1": tcrit,
    "se_ratio_cluster_over_ols": se_cl / se_ols,
    "baseline_ci_ols_low": float(b[1] - stats.t.ppf(0.975, n - k) * se_ols),
    "baseline_ci_ols_high": float(b[1] + stats.t.ppf(0.975, n - k) * se_ols),
}
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))