"""Descriptive statistics, summary tables, and exploratory figures."""
from __future__ import annotations

import numpy as np
import pandas as pd

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

set_style()

proj = pd.read_parquet(DATA / "projects.parquet")
men = pd.read_parquet(DATA / "mentions.parquet")
panel = pd.read_parquet(DATA / "panel.parquet")

TYPE_LABEL = {
    "ma_integration": "M&A integration", "capex_program": "Capex program",
    "plant_construction": "Plant construction", "partnership": "Partnership",
    "cost_reduction": "Cost reduction", "geographic_expansion": "Geographic expansion",
    "infrastructure": "Infrastructure", "product_launch": "Product launch",
    "technology_deployment": "Technology deployment", "manufacturing_expansion": "Mfg. expansion",
    "rd_program": "R&D program", "ai_initiative": "AI initiative",
    "digital_transformation": "Digital transformation", "sustainability": "Sustainability",
    "supply_chain": "Supply chain", "data_center": "Data center",
    "energy_transition": "Energy transition", "cloud_migration": "Cloud migration",
    "automation": "Automation", "drug_pipeline": "Drug pipeline",
    "erp_implementation": "ERP implementation",
}

# ---- Figure 1: mentions per year by form ----------------------------------
mp = (men[(men.year >= 2010) & (men.year <= 2026)]
      .groupby(["year", "form_type"]).size().unstack(fill_value=0))
fig, ax = plt.subplots(figsize=(7.2, 4.0))
mp.plot(ax=ax, marker="o", ms=3.5, lw=1.6)
ax.set_xlabel("Filing year"); ax.set_ylabel("Project mentions")
ax.set_title("Project mentions extracted per year, by filing form")
ax.legend(title="Form")
savefig(fig, "fig_mentions_by_year.pdf")

# ---- Figure 2: project type distribution + announced capital --------------
bt = (proj.groupby("project_type")
      .agg(projects=("project_id", "count"),
           capex=("total_amount_usd", lambda s: np.nansum(s.values)))
      .sort_values("projects", ascending=True))
bt.index = [TYPE_LABEL.get(i, i) for i in bt.index]
fig, ax = plt.subplots(figsize=(7.4, 6.2))
ax.barh(bt.index, bt["projects"], color="#1b3a5b")
ax.set_xlabel("Number of resolved projects")
ax.set_title("Distribution of projects across the 21-category taxonomy")
for y, (n, cap) in enumerate(zip(bt["projects"], bt["capex"])):
    ax.text(n + 60, y, f"{n:,}", va="center", fontsize=8)
savefig(fig, "fig_type_distribution.pdf")

# ---- Figure 3: sector x type heatmap (row-normalised share) ---------------
top_types = proj["project_type"].value_counts().head(12).index
ct = (proj[proj.project_type.isin(top_types)]
      .pivot_table(index="sector", columns="project_type", values="project_id",
                   aggfunc="count", fill_value=0))
ct = ct.div(ct.sum(axis=1), axis=0)
ct = ct.rename(columns=TYPE_LABEL)
fig, ax = plt.subplots(figsize=(9.0, 5.2))
im = ax.imshow(ct.values, cmap="magma_r", aspect="auto")
ax.set_xticks(range(len(ct.columns))); ax.set_xticklabels(ct.columns, rotation=45, ha="right", fontsize=8)
ax.set_yticks(range(len(ct.index))); ax.set_yticklabels(ct.index, fontsize=8)
ax.set_title("Project-type composition by GICS sector (row-normalised share)")
fig.colorbar(im, ax=ax, shrink=0.8, label="Within-sector share")
ax.grid(False)
savefig(fig, "fig_sector_type_heatmap.pdf")

# ---- Figure 4: announced capex by year (by selected type) -----------------
pa = proj[(proj.ann_year >= 2010) & (proj.ann_year <= 2025)].copy()
pa["ann_year"] = pa["ann_year"].astype(int)
cap_year = (pa.groupby("ann_year")["total_amount_usd"]
            .apply(lambda s: np.nansum(s.values) / 1e9))
n_year = pa.groupby("ann_year")["project_id"].count()
fig, ax = plt.subplots(figsize=(7.2, 4.0))
ax.bar(cap_year.index, cap_year.values, color="#2e8b57", alpha=0.85, label="Announced capital ($B)")
ax.set_xlabel("Announcement year"); ax.set_ylabel("Announced capital (USD bn)")
ax2 = ax.twinx(); ax2.plot(n_year.index, n_year.values, color="#c44e52", marker="o", ms=3, label="# projects")
ax2.set_ylabel("Number of projects"); ax2.grid(False)
ax.set_title("Announced project capital and project count by year")
savefig(fig, "fig_capex_by_year.pdf")

# ---- Figure 5: top locations ----------------------------------------------
loc = proj["canonical_location"].replace("", np.nan).dropna().value_counts().head(15)[::-1]
fig, ax = plt.subplots(figsize=(7.0, 4.6))
ax.barh(loc.index, loc.values, color="#8172b3")
ax.set_xlabel("Number of projects"); ax.set_title("Top 15 project locations")
savefig(fig, "fig_top_locations.pdf")

# ---- Table 1: summary statistics ------------------------------------------
def stat_row(name, s):
    s = pd.to_numeric(s, errors="coerce").dropna()
    return (name, len(s), s.mean(), s.std(), s.quantile(.5), s.min(), s.max())

rows = [
    stat_row("Projects per firm", proj.groupby("cik").size()),
    stat_row("Mentions per project", proj["n_mentions"]),
    stat_row("Filings per project", proj["n_filings"]),
    stat_row("Technologies per project", proj["n_tech"]),
    stat_row("Announced amount (USD m)", proj.loc[proj.has_amount, "total_amount_usd"] / 1e6),
    stat_row("Project duration (days)", proj.loc[proj.duration_days > 0, "duration_days"]),
    stat_row("Avg. confidence", proj["avg_confidence"]),
    stat_row("Projects/firm-year (panel)", panel["n_projects"]),
    stat_row("AI projects/firm-year", panel["n_ai"]),
]
lt = [r"\begin{tabular}{lrrrrrr}", r"\toprule",
      r"Variable & N & Mean & SD & Median & Min & Max \\", r"\midrule"]
for nm, n, mu, sd, md, mn, mx in rows:
    lt.append(f"{esc(nm)} & {n:,} & {mu:,.2f} & {sd:,.2f} & {md:,.2f} & {mn:,.2f} & {mx:,.2f} \\\\")
lt += [r"\bottomrule", r"\end{tabular}"]
write_table("tab_summary_stats", "\n".join(lt))

# ---- Table 2: projects by type --------------------------------------------
bt2 = (proj.groupby("project_type")
       .agg(n=("project_id", "count"),
            cap=("total_amount_usd", lambda s: np.nansum(s.values) / 1e9),
            ai=("is_ai", "mean"),
            conf=("avg_confidence", "mean"))
       .sort_values("n", ascending=False))
lt = [r"\begin{tabular}{lrrrr}", r"\toprule",
      r"Project type & Projects & Capital (\$bn) & AI share & Mean conf. \\", r"\midrule"]
for t, r in bt2.iterrows():
    lt.append(f"{esc(TYPE_LABEL.get(t,t))} & {int(r.n):,} & {r.cap:,.1f} & {r.ai:.2f} & {r.conf:.2f} \\\\")
lt += [r"\midrule",
       f"\\textbf{{Total}} & {len(proj):,} & {np.nansum(proj['total_amount_usd'].values)/1e9:,.1f} & "
       f"{proj['is_ai'].mean():.2f} & {proj['avg_confidence'].mean():.2f} \\\\",
       r"\bottomrule", r"\end{tabular}"]
write_table("tab_by_type", "\n".join(lt))

# ---- Table 3: sector breakdown --------------------------------------------
bs = (proj.groupby("sector")
      .agg(firms=("cik", "nunique"), n=("project_id", "count"),
           cap=("total_amount_usd", lambda s: np.nansum(s.values)/1e9),
           ai=("is_ai", "mean"))
      .sort_values("n", ascending=False))
lt = [r"\begin{tabular}{lrrrr}", r"\toprule",
      r"GICS sector & Firms & Projects & Capital (\$bn) & AI share \\", r"\midrule"]
for sct, r in bs.iterrows():
    lt.append(f"{esc(sct)} & {int(r.firms)} & {int(r.n):,} & {r.cap:,.1f} & {r.ai:.2f} \\\\")
lt += [r"\bottomrule", r"\end{tabular}"]
write_table("tab_by_sector", "\n".join(lt))

save_results("descriptives", {
    "projects_per_firm_mean": float(proj.groupby("cik").size().mean()),
    "median_amount_musd": float((proj.loc[proj.has_amount, "total_amount_usd"]/1e6).median()),
    "share_with_amount": float(proj["has_amount"].mean()),
    "top_type": bt2.index[0], "top_type_n": int(bt2.iloc[0]["n"]),
    "ai_share_overall": float(proj["is_ai"].mean()),
    "mention_cagr_2010_2025": float((men[men.year==2025].shape[0] / men[men.year==2010].shape[0]) ** (1/15) - 1),
})
print("descriptives done")