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

DATA = "data/panel.csv"
raw = pd.read_csv(DATA)
n_raw = len(raw)
issues, excl = [], []

# Exclusion steps (documented even if they remove nothing)
d = raw.copy()
n0 = len(d); d = d.dropna(subset=["firm_id", "x", "y"])
excl.append(("1", "drop rows with missing firm_id/x/y", n0, len(d), "OLS needs complete cases"))
n0 = len(d); d = d[np.isfinite(d[["x", "y"]]).all(axis=1)]
excl.append(("2", "drop rows with non-finite x/y", n0, len(d), "OLS needs finite values"))
ndup = int(d.duplicated().sum())
if ndup:
    issues.append({"issue": f"{ndup} exact duplicate rows", "handling": "kept (cannot distinguish from genuine repeats)"})
for s, r, a, b, why in excl:
    if a != b:
        issues.append({"issue": f"{a-b} rows removed by rule: {r}", "handling": why})

X = sm.add_constant(d["x"].to_numpy())
yv = d["y"].to_numpy()
g = d["firm_id"].to_numpy()
ols = sm.OLS(yv, X).fit()  # conventional homoskedastic SE (baseline)
cl = sm.OLS(yv, X).fit(cov_type="cluster", cov_kwds={"groups": g, "use_correction": True})  # CR1
G = int(len(np.unique(g)))
beta, alpha = float(cl.params[1]), float(cl.params[0])
se_c = float(cl.bse[1])
tcrit = float(stats.t.ppf(0.975, G - 1))  # small-G t(G-1) reference distribution
res = {
    "task_id": "F2-smoke-01",
    "beta_hat": beta,
    "alpha_hat": alpha,
    "se_ols": float(ols.bse[1]),
    "se_cluster": se_c,
    "beta_ci_low": beta - tcrit * se_c,
    "beta_ci_high": beta + tcrit * se_c,
    "n_analysed": int(len(d)),
    "n_clusters": G,
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary
    "t_crit_cluster": tcrit,
    "ci_df": G - 1,
    "se_ratio_cluster_over_ols": se_c / float(ols.bse[1]),
    "baseline_ci_low_ols": float(ols.conf_int()[1][0]),
    "baseline_ci_high_ols": float(ols.conf_int()[1][1]),
    "obs_per_cluster_min": int(d.groupby("firm_id").size().min()),
    "obs_per_cluster_max": int(d.groupby("firm_id").size().max()),
    "input_sha256": hashlib.sha256(open(DATA, "rb").read()).hexdigest(),
}
os.makedirs("out", exist_ok=True)
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "rationale"]).assign(
    n_removed=lambda t: t.n_before - t.n_after)[["step", "rule", "n_before", "n_after", "n_removed", "rationale"]].to_csv("out/exclusions.csv", index=False)
print(json.dumps(res, indent=2))