"""F2-dev-01: two-way FE slope with firm-clustered SE. Run from workspace root."""
import json, hashlib
import numpy as np, pandas as pd
from scipy import stats
import statsmodels.formula.api as smf

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

def step(name, rule, before, after, why):
    excl.append(dict(step=name, rule=rule, n_before=before, n_after=after,
                     n_removed=before - after, rationale=why))

d = raw.copy()
b = len(d); d = d.dropna(subset=["firm_id", "year", "x", "y"])
step(1, "drop rows with missing firm_id/year/x/y", b, len(d), "FE estimator needs complete cases")
b = len(d); d = d.drop_duplicates()
step(2, "drop exact duplicate rows", b, len(d), "duplicates carry no new information")
b = len(d); d = d[np.isfinite(d.x) & np.isfinite(d.y)]
step(3, "drop non-finite x/y", b, len(d), "non-numeric values invalid")
if d.duplicated(["firm_id", "year"]).any():
    issues.append({"issue": "conflicting duplicate firm-year keys", "handling": "kept; reported"})
balanced = d.groupby("firm_id").size().nunique() == 1 and d.groupby("year").size().nunique() == 1
if not balanced:
    issues.append({"issue": "panel unbalanced", "handling": "LSDV estimation handles imbalance"})

N = len(d); G = d.firm_id.nunique(); T = d.year.nunique()

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

# Two-way FE via LSDV (exact, handles imbalance)
fe = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit()
beta_fe = float(fe.params["x"])
se_conv = float(fe.bse["x"])
K_all = int(fe.df_model) + 1  # all estimated params incl. FE

# Firm-clustered (CR1) SE via Frisch-Waugh: residualise x on FE, use FE residuals.
fwl = smf.ols("x ~ C(firm_id) + C(year)", data=d).fit()
xt = fwl.resid.values; u = fe.resid.values
# check FWL
assert abs((xt @ d.y.values) / (xt @ xt) - beta_fe) < 1e-8
scores = pd.Series(xt * u).groupby(d.firm_id.values).sum().values
V0 = (scores @ scores) / (xt @ xt) ** 2
# Small-sample factor: firm FE are nested within clusters, so (as in Stata xtreg/reghdfe)
# they are not counted in K; K = slope + intercept + (T-1) year effects.
K_nest = 1 + 1 + (T - 1)
c = G / (G - 1) * (N - 1) / (N - K_nest)
se_cluster = float(np.sqrt(c * V0))
# Sensitivity: counting all FE parameters (statsmodels LSDV default)
c_all = G / (G - 1) * (N - 1) / (N - K_all)
se_cluster_allK = float(np.sqrt(c_all * V0))
lsdv_cl = fe.get_robustcov_results(cov_type="cluster", groups=d.firm_id.values)
se_cluster_sm = float(lsdv_cl.bse[list(fe.params.index).index("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
    "baseline_se_pooled_conventional": float(pooled.bse["x"]),
    "se_cluster_unadjusted_CR0": float(np.sqrt(V0)),
    "se_cluster_allK_dfadj": se_cluster_allK,
    "se_cluster_statsmodels_lsdv": se_cluster_sm,
    "ci_t_critical_df_G_minus_1": tcrit,
    "cluster_small_sample_factor": float(c),
    "df_resid_conventional": int(fe.df_resid),
    "panel_balanced": bool(balanced),
    "input_sha256": hashlib.sha256(open(DATA, "rb").read()).hexdigest(),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)
print(json.dumps(res, indent=2))