"""F2-smoke-01: OLS slope of y on x with firm-clustered SE. Reads data/ only; writes out/results.json."""
import json, hashlib, os
import numpy as np, pandas as pd
import statsmodels.api as sm
from scipy import stats

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
path = os.path.join(ROOT, "data", "panel.csv")
raw = pd.read_csv(path)
n_raw = len(raw)

# Exclusion steps (audit found none needed, but rules are applied and logged)
excl = []
d = raw.copy()
def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask(d)].copy(); excl.append([name, rule, nb, len(d), nb - len(d), why])
step("1", "drop rows with missing firm_id/x/y", lambda z: z[["firm_id","x","y"]].notna().all(axis=1), "OLS requires complete cases")
step("2", "drop exact duplicate rows", lambda z: ~z.duplicated(), "guard against duplicated records")
pd.DataFrame(excl, columns=["step","rule","n_before","n_after","n_removed","rationale"]).to_csv(
    os.path.join(ROOT, "out", "exclusions.csv"), index=False)

X = sm.add_constant(d["x"].to_numpy()); y = d["y"].to_numpy(); g = d["firm_id"].to_numpy()
ols = sm.OLS(y, X).fit()
cl = sm.OLS(y, X).fit(cov_type="cluster", cov_kwds={"groups": g, "use_correction": True, "df_correction": True})
G = int(len(np.unique(g))); N = len(d); K = X.shape[1]

# manual CR1 check
e = ols.resid; XtXi = np.linalg.inv(X.T @ X)
meat = sum(np.outer(X[g == k].T @ e[g == k], X[g == k].T @ e[g == k]) for k in np.unique(g))
V = (G/(G-1))*((N-1)/(N-K)) * XtXi @ meat @ XtXi
assert abs(np.sqrt(V[1,1]) - cl.bse[1]) < 1e-10

beta, alpha = float(ols.params[1]), float(ols.params[0])
se_cl = float(cl.bse[1]); tcrit = float(stats.t.ppf(0.975, G - 1))
res = {
    "task_id": "F2-smoke-01",
    "beta_hat": beta, "alpha_hat": alpha,
    "se_ols": float(ols.bse[1]), "se_cluster": se_cl,
    "beta_ci_low": beta - tcrit*se_cl, "beta_ci_high": beta + tcrit*se_cl,
    "n_analysed": int(N), "n_clusters": G,
    "feasibility": "complete", "feasibility_reason": None,
    "data_issues": [],
    "n_input_rows": {"panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary
    "se_ratio_cluster_over_ols": se_cl/float(ols.bse[1]),
    "ci_critical_value_t_df": G - 1, "ci_critical_value": tcrit,
    "baseline_ci_low_ols": beta - float(stats.t.ppf(0.975, N-K))*float(ols.bse[1]),
    "baseline_ci_high_ols": beta + float(stats.t.ppf(0.975, N-K))*float(ols.bse[1]),
    "obs_per_firm_min": int(d.groupby("firm_id").size().min()),
    "obs_per_firm_max": int(d.groupby("firm_id").size().max()),
    "input_sha256": hashlib.sha256(open(path,"rb").read()).hexdigest(),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))