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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data", "panel.csv")
OUT = os.path.join(ROOT, "out")

raw = pd.read_csv(DATA)
n_raw = len(raw)
issues, excl = [], []

def step(name, rule, df_before, df_after, why):
    excl.append(dict(step=name, rule=rule, n_before=len(df_before), n_after=len(df_after),
                     n_removed=len(df_before) - len(df_after), rationale=why))

d0 = raw
d1 = d0.dropna(subset=["firm_id", "x", "y"])
step("1", "drop rows with missing firm_id/x/y", d0, d1, "OLS requires complete cases")
if len(d1) < len(d0):
    issues.append({"issue": f"{len(d0)-len(d1)} rows with missing values", "handling": "dropped"})
d2 = d1[np.isfinite(d1["x"]) & np.isfinite(d1["y"])]
step("2", "drop rows with non-finite x/y", d1, d2, "OLS requires finite values")
if len(d2) < len(d1):
    issues.append({"issue": f"{len(d1)-len(d2)} non-finite rows", "handling": "dropped"})
d3 = d2.drop_duplicates()
step("3", "drop exact duplicate rows", d2, d3, "exact duplicates treated as data-entry errors")
if len(d3) < len(d2):
    issues.append({"issue": f"{len(d2)-len(d3)} exact duplicate rows", "handling": "dropped"})
df = d3.reset_index(drop=True)

X = sm.add_constant(df["x"].to_numpy())
y = df["y"].to_numpy()
g = df["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})
G = int(pd.Series(g).nunique())
beta, alpha = float(ols.params[1]), float(ols.params[0])
se_ols, se_cl = float(ols.bse[1]), float(cl.bse[1])
tcrit = float(stats.t.ppf(0.975, G - 1))  # t with G-1 df, standard for cluster-robust inference

# manual CR1 cross-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))
N, K = X.shape
V = (G / (G - 1)) * ((N - 1) / (N - K)) * XtXi @ meat @ XtXi
assert abs(np.sqrt(V[1, 1]) - se_cl) < 1e-10 * max(1, se_cl)

res = {
    "task_id": "F2-smoke-01",
    "beta_hat": beta,
    "alpha_hat": alpha,
    "se_ols": se_ols,
    "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": issues,
    "n_input_rows": {"panel.csv": int(n_raw)},
    "random_seed": None,
    "se_ratio_cluster_to_ols": se_cl / se_ols,
    "t_crit_cluster_df": tcrit,
    "cluster_df": G - 1,
    "baseline_ci_ols_low": beta - float(stats.t.ppf(0.975, N - K)) * se_ols,
    "baseline_ci_ols_high": beta + float(stats.t.ppf(0.975, N - K)) * se_ols,
    "obs_per_cluster_min": int(pd.Series(g).value_counts().min()),
    "obs_per_cluster_max": int(pd.Series(g).value_counts().max()),
    "input_sha256": hashlib.sha256(open(DATA, "rb").read()).hexdigest(),
}
assert n_raw - sum(r["n_removed"] for r in excl) == N
os.makedirs(OUT, exist_ok=True)
pd.DataFrame(excl).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))