"""F2-dev-01: two-way FE slope with firm-clustered SE. Run from workspace root.
No randomness is used (random_seed = null)."""
import json, hashlib
import numpy as np, pandas as pd
from scipy import stats
import statsmodels.formula.api as smf

PATH = "data/firm_panel.csv"
raw = pd.read_csv(PATH)
n_raw = len(raw)
issues = []
excl = []

d = raw.copy()
# Step 1: missing values
nb = len(d); d = d.dropna(subset=["firm_id", "year", "x", "y"])
excl.append(("1", "drop rows with missing firm_id/year/x/y", nb, len(d), "FE estimator needs complete cases"))
# Step 2: duplicate firm-year keys
nb = len(d); d = d.drop_duplicates(subset=["firm_id", "year"], keep=False)
excl.append(("2", "drop all rows of duplicated firm_id-year keys", nb, len(d), "panel must have one row per firm-year"))
if n_raw - len(d) > 0:
    issues.append({"issue": f"{n_raw-len(d)} rows missing or duplicated", "handling": "dropped"})
# Step 3: singleton firms (no within-firm variation)
nb = len(d); cnt = d.groupby("firm_id")["y"].transform("size"); d = d[cnt > 1]
excl.append(("3", "drop singleton firms", nb, len(d), "singletons carry no within-firm information"))
d = d.sort_values(["firm_id", "year"]).reset_index(drop=True)

N = len(d); G = d.firm_id.nunique(); T = d.year.nunique()
balanced = bool((d.groupby("firm_id").size() == T).all())

# Baseline: pooled OLS with intercept
pooled = smf.ols("y ~ x", data=d).fit()
beta_pooled = float(pooled.params["x"])

# Two-way FE via LSDV (exact)
lsdv = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit()
beta_fe = float(lsdv.params["x"])
se_conv = float(lsdv.bse["x"])           # homoskedastic, resid df = N - G - T
k_lsdv = int(lsdv.df_model) + 1
resid = lsdv.resid.values

# Cross-check with two-way within transformation (exact for balanced panel)
def dm(v):
    return v - v.groupby(d.firm_id).transform("mean") - v.groupby(d.year).transform("mean") + v.mean()
xt, yt = dm(d.x), dm(d.y)
beta_within = float((xt * yt).sum() / (xt ** 2).sum())

# Firm-clustered sandwich for beta via FWL (x residualised on both FE sets).
xr = smf.ols("x ~ C(firm_id) + C(year)", data=d).fit().resid.values
scores = pd.Series(xr * resid).groupby(d.firm_id.values).sum().values
v_raw = float((scores ** 2).sum()) / float((xr ** 2).sum()) ** 2
# Small-sample factor G/(G-1)*(N-1)/(N-K). Firm FE are nested within clusters and
# are not counted in K (Stata xtreg-fe / reghdfe convention): K = slope + (T-1) year effects.
K_nested = 1 + (T - 1)
c_nested = G / (G - 1) * (N - 1) / (N - K_nested)
se_cluster = float(np.sqrt(c_nested * v_raw))
# Sensitivity: all LSDV dummies counted in K (statsmodels default on LSDV)
c_lsdv = G / (G - 1) * (N - 1) / (N - k_lsdv)
se_cluster_lsdv = float(np.sqrt(c_lsdv * v_raw))
sm_cl = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit(cov_type="cluster", cov_kwds={"groups": d.firm_id.values})
se_cluster_sm = float(sm_cl.bse["x"])

tcrit = float(stats.t.ppf(0.975, G - 1))
lo, hi = beta_fe - tcrit * se_cluster, beta_fe + tcrit * se_cluster

res = {
    "task_id": "F2-dev-01",
    "beta_pooled": beta_pooled,
    "beta_fe": beta_fe,
    "se_conventional": se_conv,
    "se_cluster": se_cluster,
    "beta_ci_low": float(lo),
    "beta_ci_high": float(hi),
    "n_firms": int(G),
    "n_years": int(T),
    "n_analysed": int(N),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"firm_panel.csv": int(n_raw)},
    "random_seed": None,
    # supplementary keys
    "se_pooled_conventional": float(pooled.bse["x"]),
    "beta_fe_within_check": beta_within,
    "se_cluster_lsdv_all_dummies_counted": se_cluster_lsdv,
    "se_cluster_statsmodels_lsdv": se_cluster_sm,
    "se_ratio_cluster_to_conventional": se_cluster / se_conv,
    "cluster_small_sample_factor": c_nested,
    "ci_critical_value_t": tcrit,
    "ci_df": int(G - 1),
    "panel_balanced": balanced,
    "input_sha256": hashlib.sha256(open(PATH, "rb").read()).hexdigest(),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame([dict(step=s, rule=r, n_before=a, n_after=b, n_removed=a - b, rationale=w) for s, r, a, b, w in excl]
             ).to_csv("out/exclusions.csv", index=False)
print(json.dumps(res, indent=2))