"""Difference-in-differences event study: COVID-19 and digitalization.

Treatment = firms with above-median *physical-project intensity* in the
pre-period (2015-2019): plant construction, manufacturing/geographic expansion,
infrastructure, capex programs. Outcome = share of a firm-year's project
announcements that are digital (AI / cloud / ERP / data-center / automation /
digital-transformation / technology-deployment). We estimate dynamic
treatment x year effects (2019 omitted) with firm and year fixed effects,
testing whether disrupted incumbents accelerated digital catch-up after 2020.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

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

warnings.filterwarnings("ignore")
set_style()

proj = pd.read_parquet(DATA / "projects.parquet")
proj = proj[proj["ann_year"].between(2013, 2025)].copy()
proj["ann_year"] = proj["ann_year"].astype(int)
PHYS = ("plant_construction", "manufacturing_expansion", "geographic_expansion",
        "infrastructure", "capex_program")

g = proj.groupby(["cik", "ann_year"])
fy = g.agg(n=("project_id", "count"),
           n_digital=("is_digital", "sum"),
           n_phys=("project_type", lambda s: s.isin(PHYS).sum())).reset_index()
fy = fy.rename(columns={"ann_year": "year"})
fy["digital_share"] = fy["n_digital"] / fy["n"]

# pre-period physical intensity -> treatment
pre = fy[fy.year.between(2015, 2019)].groupby("cik").apply(
    lambda d: d["n_phys"].sum() / max(d["n"].sum(), 1)).rename("phys_pre")
med = pre.median()
treat = (pre > med).astype(int).rename("treat")
fy = fy.merge(treat, on="cik", how="inner")
fy = fy[fy["year"].between(2015, 2025)].copy()

# event-study dummies (omit 2019)
for y in sorted(fy["year"].unique()):
    if y == 2019:
        continue
    fy[f"d_{y}"] = ((fy["year"] == y) & (fy["treat"] == 1)).astype(int)

evt_terms = " + ".join([f"d_{y}" for y in sorted(fy["year"].unique()) if y != 2019])
m = smf.ols(f"digital_share ~ {evt_terms} + C(cik) + C(year)", data=fy)\
       .fit(cov_type="cluster", cov_kwds={"groups": fy["cik"]})

years = [y for y in sorted(fy["year"].unique())]
coefs, ci = [], []
for y in years:
    if y == 2019:
        coefs.append(0.0); ci.append(0.0)
    else:
        k = f"d_{y}"
        coefs.append(m.params[k]); ci.append(1.96 * m.bse[k])

fig, ax = plt.subplots(figsize=(7.4, 4.2))
ax.errorbar(years, coefs, yerr=ci, fmt="o-", color="#1b3a5b", capsize=3, ms=4)
ax.axhline(0, color="grey", lw=0.8)
ax.axvline(2019.5, color="#c44e52", ls="--", lw=1.2, label="COVID-19 onset")
ax.set_xlabel("Year"); ax.set_ylabel("Treatment $\\times$ year effect on digital share")
ax.set_title("Event study: digitalization of physically-intensive incumbents")
ax.legend()
savefig(fig, "fig_did_event.pdf")

# canonical 2x2 DiD
fy["post"] = (fy["year"] >= 2020).astype(int)
did = smf.ols("digital_share ~ treat*post + C(year)", data=fy)\
         .fit(cov_type="cluster", cov_kwds={"groups": fy["cik"]})
b = did.params["treat:post"]; se = did.bse["treat:post"]; pv = did.pvalues["treat:post"]
star = "***" if pv < .01 else "**" if pv < .05 else "*" if pv < .1 else ""

lt = [r"\begin{tabular}{lc}", r"\toprule", r" & Digital share \\", r"\midrule",
      f"Treat $\\times$ Post & {b:.4f}{star} \\\\", f" & ({se:.4f}) \\\\",
      f"Treat (pre-mean) & {did.params['treat']:.4f} \\\\",
      r"\midrule", r"Year FE & Yes \\",
      f"Observations & {int(did.nobs):,} \\\\",
      f"$R^2$ & {did.rsquared:.3f} \\\\",
      r"\bottomrule", r"\end{tabular}"]
write_table("tab_did", "\n".join(lt))

# pre-trend joint test (2015-2018 interactions == 0)
pre_terms = [f"d_{y}" for y in years if y < 2019]
ftest = m.f_test(" , ".join([f"{t} = 0" for t in pre_terms]))
save_results("did", {
    "did_coef": float(b), "did_se": float(se), "did_p": float(pv),
    "n_treat_firms": int(treat.sum()), "n_control_firms": int((treat == 0).sum()),
    "pretrend_F": float(np.ravel(ftest.fvalue)[0]), "pretrend_p": float(ftest.pvalue),
    "post_2022_effect": float(np.mean([m.params[f"d_{y}"] for y in (2022, 2023, 2024, 2025)])),
})
print(f"DiD done | treat:post={b:.4f}{star}, pretrend p={float(ftest.pvalue):.3f}")