"""Regenerate the perfmap: refit PCA, recluster, auto-derive cluster names.

Run from /Users/aabraham/hf/tabbench/. Writes:
  - perfmap_freeze.json     : new PCA basis + cluster labels
  - perfmap_regen_report.json : Spearman correlations + cluster diagnostics
"""
import json
import os
import sys
from collections import Counter, defaultdict

import numpy as np
import pandas as pd
from scipy.stats import spearmanr
from sklearn.cluster import SpectralCoclustering
from sklearn.decomposition import PCA

HERE = os.path.dirname(os.path.abspath(__file__))
os.chdir(HERE)

# Exec app.py up to gr.Blocks to get loaded per-dataset data + helpers
src = open("app.py").read()
END = src.find("with gr.Blocks(")
ns = {"__file__": os.path.join(HERE, "app.py"), "__name__": "__main__"}
exec(compile(src[:END], "app.py", "exec"), ns)

public_per_dataset = ns["public_per_dataset"]
public_per_dataset["dataset_id"] = public_per_dataset["dataset_id"].astype(str)

# Drop hidden display models (single-fit XGBoost/CatBoost/LightGBM) before
# fitting the PCA / clustering. Otherwise GBDTs get six columns per metric
# (single + ensemble) and the PCA basis is biased toward GBDT structure.
HIDDEN = ns.get("_HIDDEN_DISPLAY_MODELS", set())
if HIDDEN:
    before = public_per_dataset["model"].nunique()
    public_per_dataset = public_per_dataset[~public_per_dataset["model"].isin(HIDDEN)].copy()
    after = public_per_dataset["model"].nunique()
    print(f"Dropped hidden display models ({sorted(HIDDEN)}): {before} → {after} models in M")

# --- 1. Build the rank-of-rank matrix M (datasets x model-metric columns) ---
METRICS = ("Accuracy", "AUC", "F1_score", "Precision", "Recall", "Cross_entropy")
LOWER = {"Cross_entropy"}

rank_frames = []
for metric in METRICS:
    piv = public_per_dataset.pivot_table(
        index="dataset_id", columns="model", values=metric, aggfunc="mean",
    )
    asc = metric not in LOWER          # higher-is-better → ascending=False
    ranks = piv.rank(axis=1, ascending=not asc, method="average")
    n = piv.notna().sum(axis=1).clip(lower=1)
    rp = (ranks.sub(1, axis=0)).div(n - 1, axis=0).fillna(0.5)
    rp.columns = [f"{metric}/{m}" for m in rp.columns]
    rank_frames.append(rp)
M = pd.concat(rank_frames, axis=1)
M = 1 - M                              # higher = wins
print(f"M shape: {M.shape}  (datasets × model_metric_cols)")

# --- 2. Refit PCA(2) on M ---
pca = PCA(n_components=2, random_state=0)
scores_all = pca.fit_transform(M.values)
print(f"PCA explained variance: PC1={pca.explained_variance_ratio_[0]:.4f}  "
      f"PC2={pca.explained_variance_ratio_[1]:.4f}")

# --- 3. Spearman correlations of PC1/PC2 vs dataset metadata ---
meta = pd.DataFrame(json.load(open("public_datasets_info.json")))
meta["dataset_id"] = meta["dataset_id"].astype(str)
acc_f1 = public_per_dataset.groupby("dataset_id").agg(
    acc=("Accuracy", "mean"), f1=("F1_score", "mean"),
)
acc_f1["imbalance"] = acc_f1["acc"] - acc_f1["f1"]

scores_df = pd.DataFrame(scores_all, index=M.index, columns=["PC1", "PC2"])
joined = scores_df.merge(
    meta.set_index("dataset_id")[["rows", "features", "n_classes", "pct_cat", "avg_card"]],
    left_index=True, right_index=True, how="left",
).merge(acc_f1[["imbalance"]], left_index=True, right_index=True, how="left")

correlations = {}
for pc in ("PC1", "PC2"):
    correlations[pc] = {}
    for col in ("rows", "features", "n_classes", "pct_cat", "avg_card", "imbalance"):
        sub = joined[[pc, col]].dropna()
        if len(sub) < 5:
            correlations[pc][col] = None
        else:
            rho, p = spearmanr(sub[pc], sub[col])
            correlations[pc][col] = {"rho": float(rho), "p": float(p), "n": int(len(sub))}

print("\nSpearman ρ for PC1:")
for k, v in correlations["PC1"].items():
    if v:
        print(f"  {k:14s}  ρ={v['rho']:+.3f}  (p={v['p']:.3e})")
print("Spearman ρ for PC2:")
for k, v in correlations["PC2"].items():
    if v:
        print(f"  {k:14s}  ρ={v['rho']:+.3f}  (p={v['p']:.3e})")

# --- 4. Re-run SpectralCoclustering ---
N_CLUSTERS = 5
# clip to safe range for SVD
M_for_co = M.clip(0.0, 1.0).values
cc = SpectralCoclustering(n_clusters=N_CLUSTERS, random_state=0).fit(M_for_co)
row_labels = cc.row_labels_.astype(int)
col_labels = cc.column_labels_.astype(int)
print(f"\nCluster sizes:")
sizes = pd.Series(row_labels).value_counts().sort_index()
for c, n in sizes.items():
    print(f"  cluster {c}: {n} datasets")

# --- 5. Cluster name auto-derivation ---
# For each cluster, determine:
#   - winner family (most often #1 in per-dataset rank)
#   - 2nd family (most often #2 in rank)
#   - median rows / pct_cat → size + structure
#   - n_classes distribution → binary vs multi-class

FAMILIES = {
    "FM2":  ["seldon", "tabpfn_v3", "tabicl_v2"],
    "FM1":  ["tabicl", "tabpfn", "mitra", "limix", "tabdpt"],
    "GBDT": ["xgboost_ensemble", "catboost_ensemble", "lightgbm_ensemble",
             "xgboost", "catboost", "lightgbm"],
    "NN":   ["realmlp", "tabm", "modern_nca"],
}
FAMILY_OF = {m: fam for fam, models in FAMILIES.items() for m in models}
PRETTY = {
    "FM2":  "2nd-gen FM Supremacy",
    "FM1":  "1st-gen FM Supremacy",
    "GBDT": "GBDT Supremacy",
    "NN":   "Tuned NN Supremacy",
}
RUNNER = {
    "FM2":  "1st-gen FMs runner-up",
    "FM1":  "GBDTs runner-up",
    "GBDT": "2nd-gen FMs runner-up",
    "NN":   "FMs runner-up",
}

# Per-dataset accuracy pivot
acc_piv = public_per_dataset.pivot_table(
    index="dataset_id", columns="model", values="Accuracy", aggfunc="mean",
).dropna(axis=0, how="any")

# Per-family mean accuracy per dataset → per-dataset family ranking
fam_acc = {}
for fam, models in FAMILIES.items():
    cols = [m for m in models if m in acc_piv.columns]
    if cols:
        fam_acc[fam] = acc_piv[cols].mean(axis=1)
fam_acc = pd.DataFrame(fam_acc)
fam_rank = fam_acc.rank(axis=1, ascending=False, method="min")  # 1 = best family

def derive_cluster_name(cluster_idx, member_ids):
    sub_acc = fam_rank.loc[fam_rank.index.intersection(member_ids)]
    if sub_acc.empty:
        return f"Cluster {cluster_idx}"
    # Most common family at rank 1 and at rank 2
    rank1 = (sub_acc == 1).sum(axis=0).idxmax()
    other_cols = [c for c in sub_acc.columns if c != rank1]
    rank2_counts = (sub_acc[other_cols] == 2).sum(axis=0)
    rank2_family = rank2_counts.idxmax() if not rank2_counts.empty else None
    # Get metadata for size and structure
    meta_sub = meta[meta["dataset_id"].isin(member_ids)]
    median_rows = float(meta_sub["rows"].median()) if not meta_sub["rows"].dropna().empty else np.nan
    median_pct_cat = float(meta_sub["pct_cat"].median()) if not meta_sub["pct_cat"].dropna().empty else np.nan
    n_classes_vals = meta_sub["n_classes"].dropna().tolist()
    pct_binary = sum(1 for c in n_classes_vals if c == 2) / max(len(n_classes_vals), 1)
    # Composition descriptor
    if median_rows < 3000:
        size = "Small"
    elif median_rows < 20000:
        size = "Mid-size"
    else:
        size = "Large"
    if pct_binary > 0.8:
        cls = "binary"
    elif pct_binary < 0.3:
        cls = "multi-class"
    else:
        cls = "mixed"
    if median_pct_cat is not np.nan and median_pct_cat > 0.4:
        struct = "categorical"
    elif median_pct_cat is not np.nan and median_pct_cat < 0.1:
        struct = "numeric"
    else:
        struct = ""
    composition_parts = [size, struct, cls] if struct else [size, cls]
    composition = " ".join(p for p in composition_parts if p)
    leader_label = PRETTY.get(rank1, rank1)
    if rank2_family:
        runner_label = (
            f"{rank2_family.replace('FM2', '2nd-gen FMs').replace('FM1', '1st-gen FMs').replace('GBDT','GBDTs').replace('NN','Tuned NNs')} runner-up"
        )
    else:
        runner_label = "Mixed runner-ups"
    return f"{leader_label}<br>{runner_label}<br>{composition}"

cluster_names = {}
cluster_diag = {}
for c in sorted(set(row_labels)):
    c = int(c)  # python int for JSON
    members = M.index[row_labels == c].tolist()
    name = derive_cluster_name(c, members)
    cluster_names[c] = name
    meta_sub = meta[meta["dataset_id"].isin(members)]
    cluster_diag[c] = {
        "n_datasets": int((row_labels == c).sum()),
        "name": name,
        "median_rows": float(meta_sub["rows"].median()) if not meta_sub["rows"].dropna().empty else None,
        "median_features": float(meta_sub["features"].median()) if not meta_sub["features"].dropna().empty else None,
        "median_n_classes": float(meta_sub["n_classes"].median()) if not meta_sub["n_classes"].dropna().empty else None,
        "median_pct_cat": float(meta_sub["pct_cat"].median()) if not meta_sub["pct_cat"].dropna().empty else None,
        "pct_binary": float(sum(1 for v in meta_sub["n_classes"].dropna().tolist() if v == 2) / max(len(meta_sub["n_classes"].dropna()), 1)),
    }
    print(f"\nCluster {c} (n={cluster_diag[c]['n_datasets']}):")
    print(f"  name: {name}")
    print(f"  median rows={cluster_diag[c]['median_rows']}  features={cluster_diag[c]['median_features']}  classes={cluster_diag[c]['median_n_classes']}")
    print(f"  pct_cat median={cluster_diag[c]['median_pct_cat']}  binary={cluster_diag[c]['pct_binary']:.2f}")

# --- 6. Write the new freeze ---
freeze = {
    "pca_components": pca.components_.tolist(),
    "pca_mean": pca.mean_.tolist(),
    "feature_cols": list(M.columns),
    "dataset_ids": list(M.index),
    "row_labels": [int(x) for x in row_labels],
    "col_labels": [int(x) for x in col_labels],
    "cluster_names": {str(k): v for k, v in cluster_names.items()},
    "regenerated_at": "2026-05-29",
    "n_datasets": int(M.shape[0]),
    "n_features": int(M.shape[1]),
}
# Backup old freeze
old_freeze_path = "perfmap_freeze.json"
if os.path.exists(old_freeze_path):
    bak = old_freeze_path + ".bak.20260529"
    if not os.path.exists(bak):
        import shutil
        shutil.copyfile(old_freeze_path, bak)
        print(f"\nBacked up old freeze → {bak}")

with open(old_freeze_path, "w") as f:
    json.dump(freeze, f)
print(f"Wrote {old_freeze_path}")

# --- 7. Diagnostics report ---
report = {
    "pca_explained_variance": {
        "PC1": float(pca.explained_variance_ratio_[0]),
        "PC2": float(pca.explained_variance_ratio_[1]),
    },
    "spearman_correlations": correlations,
    "cluster_diagnostics": {str(k): v for k, v in cluster_diag.items()},
    "cluster_names_for_app": {str(k): v for k, v in cluster_names.items()},
}
with open("perfmap_regen_report.json", "w") as f:
    json.dump(report, f, indent=2)
print("Wrote perfmap_regen_report.json")