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

os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/firm_panel.csv")
n_raw = len(raw)
issues, excl = [], []
d = raw.copy()

def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[~mask].copy()
    excl.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/year/x/y", d[["firm_id","year","x","y"]].isna().any(axis=1), "cannot be used in regression")
step("2", "drop exact duplicate rows", d.duplicated(), "duplicate records")
dupk = d.duplicated(["firm_id","year"], keep=False)
if dupk.any():
    issues.append({"issue": f"{int(dupk.sum())} rows with conflicting duplicate firm-year keys", "handling": "dropped"})
step("3", "drop conflicting duplicate firm-year keys", dupk, "ambiguous observation")
for e in excl:
    if e["n_removed"] and e["step"] != "3":
        issues.append({"issue": f"{e['n_removed']} rows removed by rule: {e['rule']}", "handling": "dropped"})
d = d.sort_values(["firm_id","year"]).reset_index(drop=True)
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)

# Baseline: pooled OLS
pooled = smf.ols("y ~ x", data=d).fit()
# TWFE via LSDV
fe = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit()
G = d.firm_id.nunique(); T = d.year.nunique(); N = len(d)
# Cluster-robust (CR1, small-sample factor G/(G-1)*(N-1)/(N-K), K counts all LSDV params)
fec = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit(
    cov_type="cluster", cov_kwds={"groups": d.firm_id.values, "use_correction": True, "df_correction": True})
b = float(fe.params["x"]); se_conv = float(fe.bse["x"]); se_cl = float(fec.bse["x"])
# Sensitivity: CR1 not counting firm dummies (nested in clusters; xtreg-fe convention, K = 1 + (T-1) + 1 for x, year dummies, const)
fec0 = smf.ols("y ~ x + C(firm_id) + C(year)", data=d).fit(cov_type="cluster", cov_kwds={"groups": d.firm_id.values, "use_correction": False})
K_nf = 1 + (T - 1) + 1
se_cl_nested = float(fec0.bse["x"]) * np.sqrt(G/(G-1) * (N-1)/(N-K_nf))
tcrit = float(stats.t.ppf(0.975, G - 1))

# Cross-check by within-transformation (balanced panel => two-way demeaning exact)
def dm(v):
    return v - d.groupby("firm_id")[v.name].transform("mean") - d.groupby("year")[v.name].transform("mean") + v.mean()
xt, yt = dm(d.x), dm(d.y)
b_w = float((xt*yt).sum()/(xt**2).sum())
assert abs(b_w - b) < 1e-8, (b_w, b)

res = {
    "task_id": "F2-dev-01",
    "beta_pooled": float(pooled.params["x"]),
    "beta_fe": b,
    "se_conventional": se_conv,
    "se_cluster": se_cl,
    "beta_ci_low": b - tcrit*se_cl,
    "beta_ci_high": b + tcrit*se_cl,
    "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,
    "baseline_se_pooled_conventional": float(pooled.bse["x"]),
    "ci_critical_value_t_Gminus1": tcrit,
    "se_ratio_cluster_to_conventional": se_cl/se_conv,
    "beta_fe_within_check": b_w,
    "se_cluster_sensitivity_firmFE_not_counted": se_cl_nested,
    "se_cluster_CR0_no_correction": float(fec0.bse["x"]),
    "n_params_lsdv": int(len(fe.params)),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))