"""Determinants of announced project capital.

DV: log(total_amount_usd) for the 9,321 projects with a stated dollar figure.
Hedonic-style regressions with project, status, sector, year and firm fixed
effects; cluster-robust SE by firm.
"""
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()

p = pd.read_parquet(DATA / "projects.parquet")
d = p[p["has_amount"]].copy()
d["ann_year"] = d["ann_year"].astype(int)
d["sector"] = d["sector"].fillna("Unknown")
d["log_nfil"] = np.log(d["n_filings"].clip(lower=1))
d["is_ai_i"] = d["is_ai"].astype(int)
# baseline type = the modal type so coefficients read as deviations
base_type = d["project_type"].mode().iloc[0]
d["project_type"] = pd.Categorical(d["project_type"],
                                   categories=[base_type] + [t for t in d["project_type"].unique() if t != base_type])
groups = d["cik"]

# Model 1: sector + year FE
m1 = smf.ols("log_amount ~ is_ai_i + n_tech + log_nfil + C(status) + C(project_type) + C(sector) + C(ann_year)",
             data=d).fit(cov_type="cluster", cov_kwds={"groups": groups})
# Model 2: firm + year FE (absorb 500 firm dummies)
m2 = smf.ols("log_amount ~ is_ai_i + n_tech + log_nfil + C(status) + C(project_type) + C(cik) + C(ann_year)",
             data=d).fit(cov_type="cluster", cov_kwds={"groups": groups})

KEEP = ["is_ai_i", "n_tech", "log_nfil",
        "C(status)[T.in_progress]", "C(status)[T.completed]", "C(status)[T.planned]"]
LBL = {"is_ai_i": "AI project", "n_tech": "\\# technologies", "log_nfil": "log(\\# filings)",
       "C(status)[T.in_progress]": "Status: in progress", "C(status)[T.completed]": "Status: completed",
       "C(status)[T.planned]": "Status: planned"}


def fmt(m, k):
    if k not in m.params.index:
        return None
    b, se, pv = m.params[k], m.bse[k], m.pvalues[k]
    star = "***" if pv < .01 else "**" if pv < .05 else "*" if pv < .1 else ""
    return f"{b:.3f}{star}", f"({se:.3f})"


lt = [r"\begin{tabular}{lcc}", r"\toprule",
      r" & (1) Sector+Year FE & (2) Firm+Year FE \\", r"\midrule"]
for k in KEEP:
    a, b = fmt(m1, k), fmt(m2, k)
    lt.append(f"{LBL[k]} & {a[0] if a else ''} & {b[0] if b else ''} \\\\")
    lt.append(f" & {a[1] if a else ''} & {b[1] if b else ''} \\\\")
lt += [r"\midrule",
       r"Project-type FE & Yes & Yes \\",
       r"Sector FE & Yes & -- \\",
       r"Firm FE & -- & Yes \\",
       r"Year FE & Yes & Yes \\",
       f"Observations & {int(m1.nobs):,} & {int(m2.nobs):,} \\\\",
       f"$R^2$ & {m1.rsquared:.3f} & {m2.rsquared:.3f} \\\\",
       r"\bottomrule", r"\end{tabular}"]
write_table("tab_amount_models", "\n".join(lt))

# ---- coefficient plot: project-type premia (model 2) ----------------------
TYPE_LABEL = {
    "ma_integration": "M&A integration", "capex_program": "Capex program",
    "plant_construction": "Plant construction", "infrastructure": "Infrastructure",
    "data_center": "Data center", "energy_transition": "Energy transition",
    "product_launch": "Product launch", "rd_program": "R&D program",
    "ai_initiative": "AI initiative", "sustainability": "Sustainability",
    "cost_reduction": "Cost reduction", "geographic_expansion": "Geo. expansion",
    "manufacturing_expansion": "Mfg. expansion", "partnership": "Partnership",
    "supply_chain": "Supply chain", "automation": "Automation",
    "cloud_migration": "Cloud migration", "digital_transformation": "Digital transf.",
    "erp_implementation": "ERP", "technology_deployment": "Tech. deployment",
    "drug_pipeline": "Drug pipeline",
}
rows = []
for k in m2.params.index:
    if k.startswith("C(project_type)[T."):
        t = k.split("[T.")[1].rstrip("]")
        rows.append((TYPE_LABEL.get(t, t), m2.params[k], m2.bse[k]))
rows.sort(key=lambda r: r[1])
fig, ax = plt.subplots(figsize=(6.8, 6.2))
ys = range(len(rows))
ax.errorbar([r[1] for r in rows], list(ys),
            xerr=[1.96 * r[2] for r in rows], fmt="o", color="#1b3a5b", capsize=3, ms=4)
ax.axvline(0, color="grey", ls="--", lw=0.8)
ax.set_yticks(list(ys)); ax.set_yticklabels([r[0] for r in rows], fontsize=8)
ax.set_xlabel(f"log-capital premium vs. {TYPE_LABEL.get(base_type, base_type)} (firm+year FE)")
ax.set_title("Announced-capital premia by project type")
savefig(fig, "fig_amount_coef.pdf")

save_results("amount_models", {
    "n": int(m1.nobs), "baseline_type": base_type,
    "r2_sector": float(m1.rsquared), "r2_firm": float(m2.rsquared),
    "ai_coef": [float(m2.params["is_ai_i"]), float(m2.bse["is_ai_i"]), float(m2.pvalues["is_ai_i"])],
    "ntech_coef": [float(m2.params["n_tech"]), float(m2.pvalues["n_tech"])],
    "ai_pct_effect": float(np.exp(m2.params["is_ai_i"]) - 1),
})
print("amount models done | R2(firm FE) =", round(m2.rsquared, 3))