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

d0 = pd.read_csv("data/firm_panel.csv")
n_input = len(d0)
issues = []
d = d0.copy()
na = d[["firm_id", "year", "x", "y"]].isna().any(axis=1)
if na.any():
    issues.append({"issue": f"{int(na.sum())} rows with missing values", "handling": "dropped"})
    d = d[~na]
dup = d.duplicated(["firm_id", "year"], keep=False)
if dup.any():
    issues.append({"issue": f"{int(dup.sum())} rows with duplicated firm-year keys", "handling": "exact duplicates dropped"})
    d = d.drop_duplicates()
sizes = d.groupby("firm_id").size()
balanced = sizes.nunique() == 1 and sizes.iloc[0] == d.year.nunique()
if not balanced:
    issues.append({"issue": "panel is unbalanced", "handling": "kept all observations; FE estimator valid for unbalanced panels"})
within_sd = d.groupby("firm_id").x.std()
if (within_sd == 0).any():
    issues.append({"issue": f"{int((within_sd == 0).sum())} firms with no within variation in x", "handling": "retained (contribute nothing to beta)"})

d = d.sort_values(["firm_id", "year"]).reset_index(drop=True)
pooled = smf.ols("y ~ x", data=d).fit()
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) with statsmodels default small-sample factor G/(G-1)*(N-1)/(N-K), K incl. all dummies
fe_cl = 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})
beta = float(fe.params["x"])
se_conv = float(fe.bse["x"])
se_cl = float(fe_cl.bse["x"])
tcrit = stats.t.ppf(0.975, G - 1)

# Sensitivity: within-transformed regression (K=1 slope; like xtreg fe), CR1 with G/(G-1)*(N-1)/(N-1)
dm = d.copy()
for v in ["x", "y"]:
    dm[v + "_w"] = dm[v] - dm.groupby("firm_id")[v].transform("mean") - dm.groupby("year")[v].transform("mean") + dm[v].mean()
e = dm.y_w - beta * dm.x_w
sxx = (dm.x_w ** 2).sum()
score = (dm.x_w * e).groupby(dm.firm_id).sum()
v0 = (score ** 2).sum() / sxx ** 2
se_cl_nocorr = float(np.sqrt(v0))
se_cl_xtreg = float(np.sqrt(v0 * G / (G - 1) * (N - 1) / (N - 2)))

res = {
    "task_id": "F2-dev-01",
    "beta_pooled": float(pooled.params["x"]),
    "beta_fe": beta,
    "se_conventional": se_conv,
    "se_cluster": se_cl,
    "beta_ci_low": beta - tcrit * se_cl,
    "beta_ci_high": beta + 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_input)},
    "random_seed": None,
    "sensitivity": {
        "se_pooled_conventional": float(pooled.bse["x"]),
        "se_cluster_no_small_sample_correction": se_cl_nocorr,
        "se_cluster_K_excluding_absorbed_FE": se_cl_xtreg,
        "t_critical_df_G_minus_1": float(tcrit),
        "ci_normal_low": beta - 1.959964 * se_cl,
        "ci_normal_high": beta + 1.959964 * se_cl,
        "fe_resid_df": float(fe.df_resid),
    },
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps(res, indent=2))