"""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

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

# --- audit / exclusions (each step logged even if nothing removed)
d = raw.copy()
def step(name, rule, mask, why):
    global d
    nb = len(d); d = d[mask(d)].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/non-finite firm_id, year, x or y",
     lambda t: np.isfinite(t[["firm_id","year","x","y"]].astype(float)).all(axis=1), "FE estimator needs complete cases")
step("2", "drop exact duplicate rows", lambda t: ~t.duplicated(), "duplicates would double-weight observations")
step("3", "drop duplicated firm_id-year keys", lambda t: ~t.duplicated(["firm_id","year"], keep=False), "panel must have one row per firm-year")
pd.DataFrame(excl).to_csv("out/exclusions.csv", index=False)
counts = d.groupby("firm_id").size()
if counts.nunique() != 1 or counts.iloc[0] != d.year.nunique():
    issues.append({"issue": "unbalanced panel", "handling": "FE estimator applied to available obs"})
d = d.sort_values(["firm_id","year"]).reset_index(drop=True)

# --- baseline: pooled OLS with intercept
X0 = np.column_stack([np.ones(len(d)), d.x.values])
beta_pooled = float(np.linalg.lstsq(X0, d.y.values, rcond=None)[0][1])

# --- two-way FE via LSDV (firm + year dummies), exact for any panel
firm = pd.get_dummies(d.firm_id, drop_first=False).values.astype(float)
yr = pd.get_dummies(d.year, drop_first=True).values.astype(float)
X = np.column_stack([d.x.values, firm, yr])
y = d.y.values
N, K = X.shape
XtX_inv = np.linalg.inv(X.T @ X)
b = XtX_inv @ X.T @ y
u = y - X @ b
beta_fe = float(b[0])
G = d.firm_id.nunique(); T = d.year.nunique()
rank = np.linalg.matrix_rank(X)
assert rank == K

# conventional SE: s^2 with N-K dof (K counts all firm and year effects)
s2 = u @ u / (N - K)
se_conv = float(np.sqrt(s2 * XtX_inv[0, 0]))

# cluster-robust (CR1) by firm. Only the x-row of the sandwich is needed.
# Firm effects are nested within clusters, so (as in reghdfe / xtreg,fe) they
# are not counted in the small-sample K: K_cr = 1 (x) + (T-1) year dummies.
a = XtX_inv[0, :] @ X.T          # influence weights for beta_fe
sc = pd.Series(a * u).groupby(d.firm_id.values).sum().values
V0 = float(sc @ sc)
K_cr = 1 + (T - 1)
c_nested = G / (G - 1) * (N - 1) / (N - K_cr)
c_full = G / (G - 1) * (N - 1) / (N - K)
se_cluster = float(np.sqrt(c_nested * V0))
se_cluster_allK = float(np.sqrt(c_full * V0))
se_cluster_noadj = float(np.sqrt(V0))
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 (not required)
    "ci_critical_value_t_Gminus1": tcrit,
    "ci_dof": int(G - 1),
    "se_cluster_ratio_to_conventional": se_cluster / se_conv,
    "se_cluster_alt_K_all_dummies": se_cluster_allK,
    "se_cluster_alt_no_small_sample_adj": se_cluster_noadj,
    "ci_normal_low": beta_fe - 1.959963984540054 * se_cluster,
    "ci_normal_high": beta_fe + 1.959963984540054 * se_cluster,
    "n_parameters_lsdv": int(K),
    "input_sha256": hashlib.sha256(open(PATH, "rb").read()).hexdigest(),
}
with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: res[k] for k in ["beta_pooled","beta_fe","se_conventional","se_cluster","beta_ci_low","beta_ci_high"]}, indent=2))