"""Technology-adoption diffusion: AI initiatives.

(a) Firm-year logit for the probability of announcing >=1 AI project.
(b) Logistic (Bass-type) diffusion curve fit to the cumulative share of firms
    that have ever announced an AI project.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from scipy.optimize import curve_fit

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

warnings.filterwarnings("ignore")
set_style()

panel = pd.read_parquet(DATA / "panel.parquet").copy()
panel["sector"] = panel["sector"].fillna("Unknown")
panel["year_c"] = panel["year"] - panel["year"].min()
panel = panel.sort_values(["cik", "year"])
panel["projects_lag"] = panel.groupby("cik")["n_projects"].shift(1).fillna(0)

# ---- (a) logit ------------------------------------------------------------
logit = smf.logit("any_ai ~ year_c + post_covid + projects_lag + C(sector)",
                  data=panel).fit(disp=0, maxiter=200)
mfx = logit.get_margeff()
me = pd.Series(mfx.margeff, index=logit.params.index[1:])  # margeff excludes const
me_se = pd.Series(mfx.margeff_se, index=logit.params.index[1:])

KEEP = ["year_c", "post_covid", "projects_lag"]
LBL = {"year_c": "Year trend", "post_covid": "Post-COVID (2020+)", "projects_lag": "Projects$_{t-1}$"}
lt = [r"\begin{tabular}{lcc}", r"\toprule",
      r" & Coefficient & Avg. marginal effect \\", r"\midrule"]
for k in KEEP:
    b, se, pv = logit.params[k], logit.bse[k], logit.pvalues[k]
    star = "***" if pv < .01 else "**" if pv < .05 else "*" if pv < .1 else ""
    m_ = me[k]; m_se = me_se[k]
    lt.append(f"{LBL[k]} & {b:.3f}{star} & {m_:.4f} \\\\")
    lt.append(f" & ({se:.3f}) & ({m_se:.4f}) \\\\")
lt += [r"\midrule", r"Sector FE & Yes & \\",
       f"Observations & {int(logit.nobs):,} & \\\\",
       f"Pseudo $R^2$ & {logit.prsquared:.3f} & \\\\",
       r"\bottomrule", r"\end{tabular}"]
write_table("tab_ai_logit", "\n".join(lt))

# ---- (b) diffusion S-curve ------------------------------------------------
p = pd.read_parquet(DATA / "projects.parquet")
ai = p[p["is_ai"] & (p["ann_year"].between(2010, 2025))].copy()
ai["ann_year"] = ai["ann_year"].astype(int)
first_ai = ai.groupby("cik")["ann_year"].min()
years = np.arange(2010, 2026)
n_firms = p["cik"].nunique()
cum_adopt = np.array([(first_ai <= y).sum() for y in years]) / n_firms

def logistic(t, K, r, t0):
    return K / (1 + np.exp(-r * (t - t0)))

t = years - 2010
try:
    popt, _ = curve_fit(logistic, t, cum_adopt, p0=[0.8, 0.4, 8], maxfev=10000)
except Exception:
    popt = [cum_adopt.max(), 0.3, 8]
K, r, t0 = popt
fit = logistic(t, *popt)

# AI share of all projects per year
share = (p[p.ann_year.between(2010, 2025)].assign(y=lambda x: x.ann_year.astype(int))
         .groupby("y")["is_ai"].mean())

fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.0))
ax = axes[0]
ax.plot(years, cum_adopt, "o", color="#c44e52", label="Observed", ms=5)
ax.plot(years, fit, "-", color="#1b3a5b", lw=2, label="Logistic fit")
ax.axvline(2010 + t0, color="grey", ls=":", lw=1)
ax.set_title("Cumulative AI-project adoption across firms")
ax.set_xlabel("Year"); ax.set_ylabel("Share of firms ever announcing AI")
ax.legend()
ax.text(0.04, 0.92, f"$K$={K:.2f}, $r$={r:.2f}, $t_0$={2010+t0:.1f}",
        transform=ax.transAxes, fontsize=9, va="top")
ax2 = axes[1]
ax2.plot(share.index, share.values, "o-", color="#2e8b57", lw=1.8)
ax2.set_title("AI share of all new projects per year")
ax2.set_xlabel("Year"); ax2.set_ylabel("AI share of announcements")
savefig(fig, "fig_ai_diffusion.pdf")

save_results("diffusion", {
    "logit_post_covid": [float(logit.params["post_covid"]), float(logit.pvalues["post_covid"])],
    "mfx_year_trend": float(me["year_c"]), "mfx_post_covid": float(me["post_covid"]),
    "logistic_K": float(K), "logistic_r": float(r), "inflection_year": float(2010 + t0),
    "ai_share_2025": float(share.get(2025, np.nan)), "ai_share_2015": float(share.get(2015, np.nan)),
    "pseudo_r2": float(logit.prsquared),
})
print(f"diffusion done | inflection={2010+t0:.1f}, K={K:.2f}, r={r:.2f}")