"""F2-smoke-01: OLS slope of y on x with firm-clustered SE. Reads data/panel.csv only."""
import json, os
import numpy as np, pandas as pd
from scipy import stats

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

# Exclusion ledger (rules applied in order)
ledger, d, issues = [], raw.copy(), []
def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask(d)]; ledger.append(dict(step=name, rule=rule, n_before=nb, n_after=len(d), n_removed=nb - len(d), rationale=why))
step("1", "drop rows with missing firm_id/x/y", lambda z: z[["firm_id","x","y"]].notna().all(axis=1), "OLS needs complete cases")
step("2", "drop rows with non-finite x/y", lambda z: np.isfinite(z["x"]) & np.isfinite(z["y"]), "OLS needs finite values")
step("3", "drop exact duplicate rows", lambda z: ~z.duplicated(), "guard against duplicated records")
for e in ledger:
    if e["n_removed"] > 0:
        issues.append({"issue": f"{e['n_removed']} rows failed rule '{e['rule']}'", "handling": "excluded"})
pd.DataFrame(ledger).to_csv(os.path.join(ROOT, "out", "exclusions.csv"), index=False)

d = d.sort_values(["firm_id", "x", "y"]).reset_index(drop=True)  # deterministic order
y = d["y"].to_numpy(float); X = np.column_stack([np.ones(len(d)), d["x"].to_numpy(float)])
g = d["firm_id"].to_numpy()
n, k = X.shape
XtX_inv = np.linalg.inv(X.T @ X)
b = XtX_inv @ X.T @ y
e = y - X @ b
# Conventional homoskedastic SE
s2 = e @ e / (n - k)
se_ols = float(np.sqrt(s2 * XtX_inv[1, 1]))
# Cluster-robust (CR1, Stata-style small-sample factor) SE
groups = np.unique(g); G = len(groups)
meat = np.zeros((k, k))
for gg in groups:
    m = g == gg
    sg = X[m].T @ e[m]
    meat += np.outer(sg, sg)
c = G / (G - 1) * (n - 1) / (n - k)
V = c * XtX_inv @ meat @ XtX_inv
se_cl = float(np.sqrt(V[1, 1]))
tcrit = 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,
    "ci_critical_value_t_df": int(G - 1), "ci_critical_value": float(tcrit),
    "cluster_se_correction": "CR1: G/(G-1)*(N-1)/(N-K)",
    "se_ratio_cluster_to_ols": se_cl / se_ols,
    "obs_per_cluster_min": int(d.groupby("firm_id").size().min()),
    "obs_per_cluster_max": int(d.groupby("firm_id").size().max()),
}
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))