# ============================================================
# SPID — generate figures & LaTeX tables for the showcase doc
#   ../.venv/bin/python docs/make_figures.py
# ============================================================
from __future__ import annotations

import re
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import networkx as nx

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from spid.store import Store  # noqa: E402

HERE = Path(__file__).resolve().parent
FIG = HERE / "figures"
TAB = HERE / "tables"
FIG.mkdir(exist_ok=True); TAB.mkdir(exist_ok=True)

plt.rcParams.update({
    "font.family": "serif", "font.size": 10, "axes.titlesize": 11,
    "axes.titleweight": "bold", "axes.spines.top": False, "axes.spines.right": False,
    "figure.dpi": 150, "savefig.bbox": "tight", "axes.grid": True,
    "grid.alpha": 0.25, "grid.linewidth": 0.5,
})
NAVY = "#1f3b73"; TEAL = "#2a9d8f"; AMBER = "#e9a23b"; CRIM = "#b5374b"
LABELS = {
    "ma_integration": "M&A / Integration", "plant_construction": "Plant Construction",
    "rd_program": "R&D Program", "product_launch": "Product Launch",
    "manufacturing_expansion": "Mfg. Expansion", "partnership": "Partnership",
    "technology_deployment": "Tech Deployment", "digital_transformation": "Digital Transf.",
    "infrastructure": "Infrastructure", "capex_program": "Capex Program",
    "ai_initiative": "AI Initiative", "geographic_expansion": "Geo. Expansion",
    "cost_reduction": "Cost Reduction", "sustainability": "Sustainability",
    "energy_transition": "Energy Transition", "supply_chain": "Supply Chain",
    "drug_pipeline": "Drug Pipeline", "data_center": "Data Center",
    "cloud_migration": "Cloud Migration", "automation": "Automation",
    "erp_implementation": "ERP Implementation",
}


def bn(x):  # USD -> billions
    return (x or 0) / 1e9


def fig_by_type(st):
    rows = st.con.execute("""
        SELECT project_type, COUNT(*) n, COALESCE(SUM(total_amount_usd),0) usd
        FROM projects GROUP BY 1 ORDER BY n DESC""").fetchall()
    types = [LABELS.get(r[0], r[0]) for r in rows][::-1]
    n = [r[1] for r in rows][::-1]
    usd = [bn(r[2]) for r in rows][::-1]
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(9.2, 6.2), sharey=True)
    a1.barh(types, n, color=NAVY)
    a1.set_title("Projects by type (count)"); a1.set_xlabel("number of resolved projects")
    for i, v in enumerate(n):
        a1.text(v + max(n) * 0.01, i, str(v), va="center", fontsize=8)
    a2.barh(types, usd, color=TEAL)
    a2.set_title("Associated capital (USD bn)"); a2.set_xlabel("announced USD, billions")
    a2.tick_params(labelleft=False)
    fig.tight_layout(); fig.savefig(FIG / "by_type.pdf"); plt.close(fig)


def fig_capex_year(st):
    rows = st.con.execute("""
        SELECT EXTRACT(year FROM first_seen)::INT y, COUNT(*) n,
               COALESCE(SUM(total_amount_usd),0) usd
        FROM projects WHERE first_seen IS NOT NULL GROUP BY 1 ORDER BY 1""").fetchall()
    yrs = [r[0] for r in rows]; n = [r[1] for r in rows]; usd = [bn(r[2]) for r in rows]
    fig, ax = plt.subplots(figsize=(8.6, 4.4))
    ax.bar(yrs, usd, color=AMBER, alpha=0.85, label="announced capital (USD bn)")
    ax.set_ylabel("announced capital (USD bn)"); ax.set_xlabel("year of first mention")
    ax.set_title("New projects detected per year and associated capital")
    ax2 = ax.twinx(); ax2.grid(False)
    ax2.plot(yrs, n, color=NAVY, marker="o", lw=2, label="projects started")
    ax2.set_ylabel("projects (count)", color=NAVY)
    ax2.tick_params(axis="y", colors=NAVY)
    ax.legend(loc="upper left", fontsize=8); ax2.legend(loc="upper right", fontsize=8)
    fig.tight_layout(); fig.savefig(FIG / "capex_year.pdf"); plt.close(fig)


def fig_sector_tech(st):
    sec = st.con.execute("""
        SELECT COALESCE(sector,'(unknown)'), COUNT(*) FROM projects
        GROUP BY 1 ORDER BY 2 DESC LIMIT 9""").fetchall()
    tech = st.con.execute("""
        SELECT t, COUNT(*) FROM (SELECT unnest(technologies) t FROM projects)
        WHERE t<>'' GROUP BY 1 ORDER BY 2 DESC LIMIT 10""").fetchall()
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(9.4, 4.4))
    s_lab = [r[0] for r in sec][::-1]; s_n = [r[1] for r in sec][::-1]
    a1.barh(s_lab, s_n, color=NAVY); a1.set_title("Projects by GICS sector")
    a1.set_xlabel("projects")
    t_lab = [r[0] for r in tech][::-1]; t_n = [r[1] for r in tech][::-1]
    a2.barh(t_lab, t_n, color=CRIM); a2.set_title("Projects by technology tag")
    a2.set_xlabel("projects")
    fig.tight_layout(); fig.savefig(FIG / "sector_tech.pdf"); plt.close(fig)


def fig_gantt(st):
    rows = st.con.execute("""
        SELECT ticker, project_name, first_seen, last_seen, n_filings, project_type
        FROM projects WHERE n_filings>=3 AND first_seen IS NOT NULL
        ORDER BY n_filings DESC, total_amount_usd DESC NULLS LAST LIMIT 16""").fetchall()
    rows = rows[::-1]
    fig, ax = plt.subplots(figsize=(9.2, 5.6))
    cmap = plt.get_cmap("tab20")
    for i, (tic, name, fs, ls, nf, pt) in enumerate(rows):
        start = fs.year + (fs.month - 1) / 12
        end = ls.year + (ls.month - 1) / 12
        ax.barh(i, max(end - start, 0.15), left=start, height=0.55,
                color=cmap(i % 20), alpha=0.9)
        ax.text(start - 0.15, i, f"{tic}: {name[:34]}", va="center", ha="right", fontsize=7.2)
        ax.text(end + 0.1, i, f"{nf}f", va="center", fontsize=7, color="#444")
    ax.set_yticks([]); ax.set_xlabel("year")
    ax.set_title("Project lifecycles reconstructed across filings (top multi-filing projects)")
    ax.set_xlim(2009, 2027.5); ax.margins(y=0.02)
    fig.tight_layout(); fig.savefig(FIG / "gantt.pdf"); plt.close(fig)


def fig_tesla_timeline(st):
    """Case study: extract GWh/MWh deployments from Tesla energy-storage snippets."""
    pid = st.con.execute("""
        SELECT project_id FROM projects WHERE ticker='TSLA'
        AND lower(project_name) LIKE '%energy storage%'
        ORDER BY n_filings DESC LIMIT 1""").fetchone()
    if not pid:
        return False
    tl = st.con.execute("""
        SELECT filing_date, snippet FROM project_timeline
        WHERE project_id=? ORDER BY filing_date""", [pid[0]]).fetchall()
    pts = []
    for d, snip in tl:
        m = re.search(r"([\d.]+)\s*(GWh|MWh)", snip, re.I)
        if m:
            v = float(m.group(1))
            if m.group(2).lower() == "mwh":
                v /= 1000.0
            pts.append((d, v, snip))
    if len(pts) < 3:
        return False
    fig, ax = plt.subplots(figsize=(8.6, 4.2))
    xs = [p[0] for p in pts]; ys = [p[1] for p in pts]
    ax.plot(xs, ys, color=TEAL, marker="o", lw=2)
    ax.fill_between(xs, ys, color=TEAL, alpha=0.12)
    ax.set_title("Case study — Tesla energy-storage deployments tracked filing-by-filing")
    ax.set_ylabel("energy storage deployed (GWh, as reported)")
    ax.set_xlabel("filing date")
    fig.autofmt_xdate()
    fig.tight_layout(); fig.savefig(FIG / "tesla_timeline.pdf"); plt.close(fig)
    return True


def fig_kg(st):
    """Ego knowledge graph for NVDA: company -> projects -> technologies/locations."""
    cik = st.con.execute("SELECT cik FROM projects WHERE ticker='NVDA' LIMIT 1").fetchone()
    if not cik:
        return
    cik = cik[0]
    projs = st.con.execute("""
        SELECT project_id, project_name, project_type, technologies, canonical_location
        FROM projects WHERE cik=? ORDER BY n_filings DESC, total_amount_usd DESC NULLS LAST
        LIMIT 9""", [cik]).fetchall()
    G = nx.Graph()
    comp = "NVIDIA"
    G.add_node(comp, kind="company")
    for pid, name, pt, techs, loc in projs:
        pn = name[:26]
        G.add_node(pn, kind="project"); G.add_edge(comp, pn)
        for t in (techs or [])[:2]:
            G.add_node(t, kind="tech"); G.add_edge(pn, t)
        if loc:
            ln = loc.title()[:18]
            G.add_node(ln, kind="loc"); G.add_edge(pn, ln)
    pos = nx.spring_layout(G, k=0.9, seed=7, iterations=200)
    kinds = nx.get_node_attributes(G, "kind")
    style = {"company": (NAVY, 2400), "project": (TEAL, 950),
             "tech": (CRIM, 600), "loc": (AMBER, 600)}
    fig, ax = plt.subplots(figsize=(9.2, 6.4))
    nx.draw_networkx_edges(G, pos, ax=ax, edge_color="#bbb", width=1.0)
    for kind, (col, size) in style.items():
        nodes = [n for n in G if kinds.get(n) == kind]
        nx.draw_networkx_nodes(G, pos, nodelist=nodes, node_color=col,
                               node_size=size, ax=ax, alpha=0.92, edgecolors="white")
    nx.draw_networkx_labels(G, pos, font_size=7, ax=ax)
    ax.set_title("Knowledge-graph slice — NVIDIA: projects, technologies, locations")
    ax.axis("off")
    handles = [plt.Line2D([0], [0], marker="o", color="w", markerfacecolor=c,
               markersize=10, label=l) for l, (c, _) in
               zip(["Company", "Project", "Technology", "Location"], style.values())]
    ax.legend(handles=handles, loc="lower left", fontsize=8, frameon=True)
    fig.tight_layout(); fig.savefig(FIG / "kg_nvda.pdf"); plt.close(fig)


# ----------------------------- tables ---------------------------------
def _esc(s):
    return str(s).replace("&", "\\&").replace("%", "\\%").replace("_", "\\_").replace("$", "\\$")


def tab_top_ma(st):
    rows = st.con.execute("""
        SELECT ticker, project_name, status, total_amount_usd
        FROM projects WHERE project_type='ma_integration' AND total_amount_usd IS NOT NULL
        ORDER BY total_amount_usd DESC LIMIT 10""").fetchall()
    lines = [r"\begin{tabular}{lllr}", r"\toprule",
             r"Ticker & Project & Status & Value (USD bn) \\", r"\midrule"]
    for tic, name, status, usd in rows:
        lines.append(f"{_esc(tic)} & {_esc(name[:42])} & {_esc(status)} & {bn(usd):.1f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "top_ma.tex").write_text("\n".join(lines))


def tab_top_companies(st):
    rows = st.con.execute("""
        SELECT ticker, COALESCE(company_name,ticker), sector, COUNT(*) n,
               COALESCE(SUM(total_amount_usd),0) usd
        FROM projects GROUP BY 1,2,3 ORDER BY n DESC LIMIT 12""").fetchall()
    lines = [r"\begin{tabular}{lllrr}", r"\toprule",
             r"Ticker & Company & Sector & Projects & Capital (bn) \\", r"\midrule"]
    for tic, name, sec, n, usd in rows:
        lines.append(f"{_esc(tic)} & {_esc(name[:22])} & {_esc((sec or '')[:20])} & {n} & {bn(usd):.1f} \\\\")
    lines += [r"\bottomrule", r"\end{tabular}"]
    (TAB / "top_companies.tex").write_text("\n".join(lines))


def stats(st):
    o = {
        "projects": st.count("projects"),
        "mentions": st.count("project_mentions"),
        "companies": st.con.execute("SELECT COUNT(DISTINCT cik) FROM projects").fetchone()[0],
        "capex_bn": bn(st.con.execute("SELECT COALESCE(SUM(total_amount_usd),0) FROM projects").fetchone()[0]),
        "kg_nodes": st.count("kg_nodes"), "kg_edges": st.count("kg_edges"),
        "timeline": st.count("project_timeline"),
        "types": st.con.execute("SELECT COUNT(DISTINCT project_type) FROM projects").fetchone()[0],
    }
    macro = "\n".join([
        f"\\newcommand{{\\Nprojects}}{{{o['projects']}}}",
        f"\\newcommand{{\\Nmentions}}{{{o['mentions']}}}",
        f"\\newcommand{{\\Ncompanies}}{{{o['companies']}}}",
        f"\\newcommand{{\\Ncapex}}{{{o['capex_bn']:,.0f}}}",
        f"\\newcommand{{\\Nkgnodes}}{{{o['kg_nodes']}}}",
        f"\\newcommand{{\\Nkgedges}}{{{o['kg_edges']}}}",
        f"\\newcommand{{\\Ntimeline}}{{{o['timeline']}}}",
        f"\\newcommand{{\\Ntypes}}{{{o['types']}}}",
    ])
    (TAB / "stats.tex").write_text(macro)
    print("stats:", o)


def main():
    st = Store(read_only=True)
    try:
        fig_by_type(st); print("✓ by_type")
        fig_capex_year(st); print("✓ capex_year")
        fig_sector_tech(st); print("✓ sector_tech")
        fig_gantt(st); print("✓ gantt")
        ok = fig_tesla_timeline(st); print("✓ tesla_timeline" if ok else "… tesla skipped")
        fig_kg(st); print("✓ kg_nvda")
        tab_top_ma(st); tab_top_companies(st); print("✓ tables")
        stats(st)
    finally:
        st.close()
    print("Figures in", FIG)


if __name__ == "__main__":
    main()