"""Semantic structure of the project corpus via sentence embeddings.

KMeans clustering of the 384-d MiniLM project embeddings into latent themes,
silhouette diagnostics, a t-SNE map, and cluster characterization by dominant
project type / technologies / representative names.
"""
from __future__ import annotations

import warnings
from collections import Counter

import numpy as np
import pandas as pd
from sklearn.preprocessing import normalize
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.manifold import TSNE

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

warnings.filterwarnings("ignore")
set_style()
rng = np.random.default_rng(42)

emb = pd.read_parquet(DATA / "embeddings.parquet")
proj = pd.read_parquet(DATA / "projects.parquet")
m = emb.merge(proj[["project_id", "project_type", "sector", "project_name",
                    "technologies", "total_amount_usd", "is_ai"]],
              on="project_id", how="inner")
X = normalize(np.vstack(m["vector"].apply(np.asarray).values).astype(np.float32))
print("embedding matrix:", X.shape)

# ---- choose K via silhouette (on a sample) --------------------------------
samp = rng.choice(len(X), size=min(4000, len(X)), replace=False)
sils = {}
for k in (6, 8, 10, 12, 14):
    km = KMeans(n_clusters=k, n_init=4, random_state=0).fit(X[samp])
    sils[k] = silhouette_score(X[samp], km.labels_)
K = max(sils, key=sils.get)
km = KMeans(n_clusters=K, n_init=8, random_state=0).fit(X)
m["cluster"] = km.labels_

# ---- characterize clusters ------------------------------------------------
def top_techs(sub, n=3):
    c = Counter()
    for t in sub["technologies"]:
        for x in ([] if t is None else list(t)):
            if str(x).strip():
                c[str(x).strip().title()] += 1
    return ", ".join(t for t, _ in c.most_common(n))

rows = []
for cl in range(K):
    sub = m[m.cluster == cl]
    dom_type = sub["project_type"].mode().iloc[0]
    dom_sec = sub["sector"].mode().iloc[0] if sub["sector"].notna().any() else "-"
    rows.append({
        "cluster": cl + 1, "n": len(sub), "dom_type": dom_type, "dom_sector": dom_sec,
        "ai_share": sub["is_ai"].mean(), "techs": top_techs(sub),
        "capex": np.nansum(sub["total_amount_usd"].values) / 1e9,
    })
cl_df = pd.DataFrame(rows).sort_values("n", ascending=False)

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",
}
lt = [r"\begin{tabular}{rlllr}", r"\toprule",
      r"Cl. & Projects & Dominant type & Top technologies & Capital (\$bn) \\", r"\midrule"]
for _, r in cl_df.iterrows():
    lt.append(f"{int(r.cluster)} & {int(r.n):,} & {esc(TYPE_LABEL.get(r.dom_type, r.dom_type))} & "
              f"{esc(r.techs[:42])} & {r.capex:,.1f} \\\\")
lt += [r"\bottomrule", r"\end{tabular}"]
write_table("tab_clusters", "\n".join(lt))

# ---- t-SNE map (stratified sample for speed) ------------------------------
sviz = rng.choice(len(X), size=min(5000, len(X)), replace=False)
emb2 = TSNE(n_components=2, perplexity=35, init="pca", random_state=0,
            max_iter=750).fit_transform(X[sviz])
fig, ax = plt.subplots(figsize=(8.4, 6.8))
palette = ["#1b3a5b", "#c44e52", "#2e8b57", "#dd8452", "#8172b3", "#937860",
           "#da8bc3", "#64b5cd", "#ccb974", "#55a868", "#b07aa1", "#e15759",
           "#76b7b2", "#59a14f"]
labs = km.labels_[sviz]
for cl in range(K):
    pts = emb2[labs == cl]
    ax.scatter(pts[:, 0], pts[:, 1], s=7, alpha=0.55, color=palette[cl % len(palette)],
               label=f"C{cl+1}: {TYPE_LABEL.get(cl_df.set_index('cluster').loc[cl+1,'dom_type'] if (cl+1) in cl_df['cluster'].values else '', '')[:14]}")
ax.set_title(f"t-SNE map of project embeddings ({K} KMeans themes)")
ax.set_xlabel("t-SNE 1"); ax.set_ylabel("t-SNE 2"); ax.grid(False)
ax.legend(fontsize=6, markerscale=2, ncol=2, loc="upper right")
savefig(fig, "fig_embedding_tsne.pdf")

# silhouette figure
fig, ax = plt.subplots(figsize=(5.6, 3.4))
ax.plot(list(sils.keys()), list(sils.values()), "o-", color="#1b3a5b")
ax.axvline(K, color="#c44e52", ls="--", lw=1)
ax.set_xlabel("Number of clusters $K$"); ax.set_ylabel("Silhouette score")
ax.set_title("Cluster-count selection")
savefig(fig, "fig_silhouette.pdf")

save_results("embeddings", {
    "n": int(len(X)), "dim": int(X.shape[1]), "K": int(K),
    "silhouette": {int(k): float(v) for k, v in sils.items()},
    "clusters": cl_df.to_dict("records"),
})
print(f"embeddings done | K={K}, best silhouette={sils[K]:.3f}")