"""F2-smoke-01: OLS slope of y on x with firm-clustered SEs. Reads only data/panel.csv."""
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__)))
OUT = os.path.join(ROOT, "out")
df_raw = pd.read_csv(os.path.join(ROOT, "data", "panel.csv"))
n_raw = len(df_raw)

# Exclusion ledger: drop rows with missing / non-finite values or missing firm_id
excl = []
df = df_raw.copy()
mask = np.isfinite(df[["x", "y"]]).all(axis=1) & df["firm_id"].notna()
excl.append(("1", "drop rows with missing/non-finite x, y or firm_id", len(df), int(mask.sum()),
             int((~mask).sum()), "OLS requires complete finite observations"))
df = df[mask]
pd.DataFrame(excl, columns=["step", "rule", "n_before", "n_after", "n_removed", "rationale"]).to_csv(
    os.path.join(OUT, "exclusions.csv"), index=False)

data_issues = []
if (~mask).sum() > 0:
    data_issues.append({"issue": f"{int((~mask).sum())} incomplete rows", "handling": "dropped"})
if df_raw.duplicated().sum() > 0:
    data_issues.append({"issue": f"{int(df_raw.duplicated().sum())} duplicate rows", "handling": "retained"})

y = df["y"].to_numpy(float)
X = np.column_stack([np.ones(len(df)), df["x"].to_numpy(float)])
g = df["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  # residuals

# Baseline: conventional homoskedastic OLS SE
s2 = e @ e / (N - K)
se_ols = float(np.sqrt(s2 * XtX_inv[1, 1]))

# Cluster-robust (CR1, Stata-style small-sample factor) by firm_id
groups = np.unique(g)
G = len(groups)
meat = np.zeros((K, K))
for gid in groups:
    idx = g == gid
    sg = X[idx].T @ e[idx]
    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 = 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": data_issues,
    "n_input_rows": {"panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary (not required)
    "se_ratio_cluster_over_ols": se_cl / se_ols,
    "ci_t_critical_df_G_minus_1": tcrit,
    "se_cluster_alpha": float(np.sqrt(V[0, 0])),
    "baseline_ci_low_ols": float(b[1] - stats.t.ppf(0.975, N - K) * se_ols),
    "baseline_ci_high_ols": float(b[1] + stats.t.ppf(0.975, N - K) * se_ols),
    "share_between_firm_var_resid": float(pd.Series(e).groupby(g).mean().var() / pd.Series(e).var()),
    "share_between_firm_var_x": float(df.groupby("firm_id")["x"].mean().var() / df["x"].var()),
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))