"""Predictive model: which projects are 'mega-projects'?

Target: among projects with a disclosed dollar amount, those above the median
(>$718m). Features are text-derived structural attributes only (project type,
sector, status, #technologies, #filings, #mentions, AI flag, has-location,
announcement year) -- no leakage of the amount. Gradient-boosted trees with
cross-validated ROC-AUC and permutation importance.
"""
from __future__ import annotations

import warnings
import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.model_selection import cross_val_predict, StratifiedKFold
from sklearn.metrics import roc_auc_score, roc_curve, average_precision_score
from sklearn.inspection import permutation_importance

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

warnings.filterwarnings("ignore")
set_style()

p = pd.read_parquet(DATA / "projects.parquet")
d = p[p["has_amount"]].copy()
thr = d["total_amount_usd"].median()
d["mega"] = (d["total_amount_usd"] > thr).astype(int)
d["has_loc"] = (d["canonical_location"].fillna("") != "").astype(int)
d["ann_year"] = d["ann_year"].astype(int)

feat_num = ["n_tech", "n_filings", "n_mentions", "is_ai", "has_loc", "ann_year", "avg_confidence"]
d["is_ai"] = d["is_ai"].astype(int)
Xcat = pd.get_dummies(d[["project_type", "sector", "status"]].astype(str), dtype=float)
X = pd.concat([d[feat_num].reset_index(drop=True), Xcat.reset_index(drop=True)], axis=1)
y = d["mega"].values

clf = HistGradientBoostingClassifier(max_iter=350, learning_rate=0.06,
                                     max_leaf_nodes=31, l2_regularization=1.0,
                                     random_state=0)
cv = StratifiedKFold(5, shuffle=True, random_state=0)
proba = cross_val_predict(clf, X, y, cv=cv, method="predict_proba")[:, 1]
auc = roc_auc_score(y, proba)
ap = average_precision_score(y, proba)

# fit on full for importance
clf.fit(X, y)
pi = permutation_importance(clf, X, y, n_repeats=8, random_state=0, scoring="roc_auc")
imp = pd.Series(pi.importances_mean, index=X.columns).sort_values(ascending=False)
# collapse one-hot back to groups
def group(col):
    for g in ("project_type_", "sector_", "status_"):
        if col.startswith(g):
            return g.rstrip("_")
    return col
gimp = imp.groupby(group).sum().sort_values(ascending=False)

# ROC + importance figure
fpr, tpr, _ = roc_curve(y, proba)
fig, axes = plt.subplots(1, 2, figsize=(11.0, 4.2))
ax = axes[0]
ax.plot(fpr, tpr, color="#1b3a5b", lw=2, label=f"GBT (AUC={auc:.3f})")
ax.plot([0, 1], [0, 1], "--", color="grey", lw=0.9)
ax.set_xlabel("False positive rate"); ax.set_ylabel("True positive rate")
ax.set_title("Mega-project classifier (5-fold CV ROC)"); ax.legend(loc="lower right")
ax2 = axes[1]
gg = gimp.head(8)[::-1]
LBL = {"project_type": "Project type", "sector": "Sector", "status": "Status",
       "n_tech": "\\# technologies", "n_filings": "\\# filings", "n_mentions": "\\# mentions",
       "is_ai": "AI flag", "has_loc": "Has location", "ann_year": "Year", "avg_confidence": "Confidence"}
ax2.barh([LBL.get(i, i).replace("\\#", "#") for i in gg.index], gg.values, color="#c44e52")
ax2.set_xlabel("Permutation importance (AUC drop)")
ax2.set_title("Feature importance")
savefig(fig, "fig_ml_roc.pdf")

lt = [r"\begin{tabular}{lr}", r"\toprule", r"Feature group & Importance (AUC drop) \\", r"\midrule"]
for i, v in gimp.head(8).items():
    lt.append(f"{LBL.get(i, esc(i))} & {v:.4f} \\\\")
lt += [r"\midrule", f"CV ROC-AUC & {auc:.3f} \\\\", f"Avg. precision & {ap:.3f} \\\\",
       f"Observations & {len(d):,} \\\\", r"\bottomrule", r"\end{tabular}"]
write_table("tab_ml", "\n".join(lt))

save_results("ml", {
    "n": int(len(d)), "threshold_usd": float(thr), "auc": float(auc), "avg_precision": float(ap),
    "top_features": {str(k): float(v) for k, v in gimp.head(6).items()},
})
print(f"ML done | AUC={auc:.3f}, AP={ap:.3f}")