"""Panel count models for project-announcement intensity.

DV: n_projects (count of strategic projects firm i announces in year t).
Estimators: pooled Poisson (sector+year FE), Poisson with firm+year FE (PPML),
Negative Binomial (overdispersion). Cluster-robust SE by firm.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from scipy import stats

from common import DATA, set_style, savefig, write_table, save_results
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")
set_style()

df = pd.read_parquet(DATA / "panel.parquet").copy()
df = df.dropna(subset=["n_projects_lag"]).copy()   # drop first year per firm (lag)
df["projects_lag"] = df["n_projects_lag"]
df["cum_ai_lag"] = df.groupby("cik")["cum_ai"].shift(1).fillna(0)
df["year_c"] = df["year"] - df["year"].min()
df["sector"] = df["sector"].fillna("Unknown")
groups = df["cik"]

# ---- Model 1: pooled Poisson, sector + year FE ----------------------------
f1 = "n_projects ~ projects_lag + cum_ai_lag + post_covid + year_c + C(sector)"
m1 = smf.poisson(f1, data=df).fit(cov_type="cluster", cov_kwds={"groups": groups}, disp=0, maxiter=200)

# ---- Model 2: Poisson with firm + year FE (PPML) --------------------------
f2 = "n_projects ~ projects_lag + cum_ai_lag + C(cik) + C(year)"
m2 = smf.poisson(f2, data=df).fit(cov_type="cluster", cov_kwds={"groups": groups}, disp=0, maxiter=300)

# ---- Model 3: Negative Binomial (sector + year FE) ------------------------
f3 = "n_projects ~ projects_lag + cum_ai_lag + post_covid + year_c + C(sector)"
m3 = smf.negativebinomial(f3, data=df).fit(cov_type="cluster", cov_kwds={"groups": groups}, disp=0, maxiter=300)

# overdispersion test (Cameron-Trivedi style): regress aux on fitted
mu = m1.predict(df)
aux = ((df["n_projects"] - mu) ** 2 - df["n_projects"]) / mu
od = sm.OLS(aux, mu).fit()
od_t = float(od.tvalues.iloc[0]); od_p = float(od.pvalues.iloc[0])

KEEP = ["projects_lag", "cum_ai_lag", "post_covid", "year_c"]
LBL = {"projects_lag": "Projects$_{t-1}$", "cum_ai_lag": "Cumulative AI$_{t-1}$",
       "post_covid": "Post-COVID (2020+)", "year_c": "Year trend"}


def col(m, label):
    out = {"_label": label, "_n": int(m.nobs), "_ll": float(m.llf)}
    for k in KEEP:
        if k in m.params.index:
            b, se, p = m.params[k], m.bse[k], m.pvalues[k]
            star = "***" if p < .01 else "**" if p < .05 else "*" if p < .1 else ""
            out[k] = (f"{b:.3f}{star}", f"({se:.3f})")
    return out


c1, c2, c3 = col(m1, "Pooled Poisson"), col(m2, "Poisson FE"), col(m3, "Neg. Binomial")
cols = [c1, c2, c3]

lt = [r"\begin{tabular}{lccc}", r"\toprule",
      r" & (1) & (2) & (3) \\",
      r" & Pooled Poisson & Poisson (firm+yr FE) & Neg. Binomial \\", r"\midrule"]
for k in KEEP:
    line = LBL[k]
    se_line = ""
    for cc in cols:
        if k in cc:
            line += f" & {cc[k][0]}"; se_line += f" & {cc[k][1]}"
        else:
            line += " & "; se_line += " & "
    lt.append(line + r" \\"); lt.append(se_line + r" \\")
lt.append(r"\midrule")
lt.append(r"Sector FE & Yes & -- & Yes \\")
lt.append(r"Firm FE & -- & Yes & -- \\")
lt.append(r"Year FE & Yes & Yes & Yes \\")
lt.append("Observations" + "".join(f" & {cc['_n']:,}" for cc in cols) + r" \\")
lt.append("Log-likelihood" + "".join(f" & {cc['_ll']:,.0f}" for cc in cols) + r" \\")
lt += [r"\bottomrule", r"\end{tabular}"]
write_table("tab_count_models", "\n".join(lt))

# ---- coefficient plot (Poisson FE incidence-rate effects) -----------------
fig, ax = plt.subplots(figsize=(6.4, 3.4))
ys, labels, los, his, bs = [], [], [], [], []
for i, k in enumerate(["projects_lag", "cum_ai_lag"]):
    b, se = m2.params[k], m2.bse[k]
    bs.append(b); los.append(b - 1.96 * se); his.append(b + 1.96 * se)
    labels.append(LBL[k]); ys.append(i)
ax.errorbar(bs, ys, xerr=[np.array(bs) - np.array(los), np.array(his) - np.array(bs)],
            fmt="o", color="#1b3a5b", capsize=4)
ax.axvline(0, color="grey", lw=0.8, ls="--")
ax.set_yticks(ys); ax.set_yticklabels(labels)
ax.set_xlabel("Poisson coefficient (firm + year FE)")
ax.set_title("Drivers of project-announcement intensity")
savefig(fig, "fig_count_coef.pdf")

save_results("count_models", {
    "overdispersion_t": od_t, "overdispersion_p": od_p,
    "poisson_fe": {k: [float(m2.params[k]), float(m2.bse[k]), float(m2.pvalues[k])]
                   for k in ["projects_lag", "cum_ai_lag"]},
    "nb_alpha": float(m3.params.get("alpha", np.nan)),
    "persistence_irr": float(np.exp(m2.params["projects_lag"])),
    "ai_irr": float(np.exp(m2.params["cum_ai_lag"])),
    "n_obs": int(m1.nobs),
})
print("count models done | overdispersion t =", round(od_t, 2))