# %% [markdown]
# # 3 - LightGBM and CatBoost: the two classical models
#
# The organisers ask for at least two models, chosen on Madrid. These are them.
# They are also the honest yardstick the fine-tuned LLM has to beat, and the
# source of the probabilities we later ensemble with it.
#
# Running order:
#
# 1. **Spatial** cross-validation on Madrid - random folds are dishonest here,
#    because neighbouring pixels share a construction era.
# 2. **Adversarial validation** - find the features that only say "this is
#    Madrid" and drop them.
# 3. **Per-city normalisation** - the single biggest transfer win, and it needs
#    no labels.
# 4. **LightGBM vs CatBoost vs ordinal LightGBM**, all judged on Madrid only.
# 5. **The official few-shot protocol** at n = 5, 25, 50, 100, 200.
#
# Runs on CPU. About 20 minutes on a laptop.

# %%
import sys
from pathlib import Path

for cand in [Path.cwd() / "src", Path.cwd().parent / "src",
             Path("/kaggle/working/src"), Path("/content/hack4dev/src")]:
    if cand.exists():
        sys.path.insert(0, str(cand))
        break

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import config as C
import evaluation as E
import gbdt as G
import retrieval as R

RNG = np.random.default_rng(C.RANDOM_SEED)

# Prefer the full feature tables from notebook 1. Fall back to the organisers'
# preprocessed_data.pkl, which carries the official 60 features and the pixel
# coordinates - enough for everything in THIS notebook, and only ~50 MB, so it
# is the right thing to grab first if the raw parquet is hard to get hold of.
import features as F

if (C.PROC_DIR / "madrid_features.parquet").exists():
    madrid = pd.read_parquet(C.PROC_DIR / "madrid_features.parquet")
    amsterdam = pd.read_parquet(C.PROC_DIR / "amsterdam_features.parquet")
    SOURCE = "notebook 1 (official 60 + extended)"
else:
    pkl = None
    for root in [C.DATA_DIR, C.PROJECT_ROOT, Path("/kaggle/input")]:
        if root.exists():
            hits = sorted(root.rglob("preprocessed_data.pkl"))
            if hits:
                pkl = hits[0]
                break
    if pkl is None:
        raise FileNotFoundError(
            "Neither data/processed/madrid_features.parquet (run notebook 1) "
            "nor preprocessed_data.pkl were found.")
    madrid, amsterdam = F.from_official_pkl(pkl)
    SOURCE = "organisers' preprocessed_data.pkl (official 60 only)"

print(f"\nfeature source: {SOURCE}")
print(madrid.shape, amsterdam.shape)

# %%
# Feature sets we will compare. Keeping them explicit makes the ablation table
# in the report trivial to produce.
DROP = {"px_key", "py_key", "age_class", "weighted_mean_year", "city"}
ALL_FEATS = [c for c in madrid.columns
             if c not in DROP and madrid[c].dtype.kind in "fiu"]
DEAD = [c for c in ALL_FEATS if madrid[c].std() == 0 or amsterdam[c].std() == 0]
ALL_FEATS = [c for c in ALL_FEATS if c not in DEAD]

FEATURE_SETS = {
    "official60": [c for c in C.FEATURE_COLS_60 if c in ALL_FEATS],
    "reflectance": [c for c in ALL_FEATS if c.startswith("sr_")],
    "events": [c for c in ALL_FEATS if c.startswith("step_")],
    "neighbourhood": [c for c in ALL_FEATS if c.startswith("nbr_")],
    "everything": ALL_FEATS,
}
print("dead (zero variance in at least one city):", DEAD)
for k, v in FEATURE_SETS.items():
    print(f"  {k:<14} {len(v)} features")

y_mad = madrid["age_class"].to_numpy()
y_ams = amsterdam["age_class"].to_numpy()          # opened only for scoring
coords_mad = madrid[C.PID].to_numpy(float)
coords_ams = amsterdam[C.PID].to_numpy(float)

# %% [markdown]
# ## 3.1 An honest validation split
#
# `RepeatedStratifiedKFold` on pixels puts a pixel's own neighbours in the
# training fold. Neighbours share an era, so the model can almost look up the
# answer, and Madrid CV comes out flattering - part of why 0.63 on Madrid turns
# into 0.34 on Amsterdam.
#
# Instead: cut Madrid into geographic blocks and hold out whole blocks. That is
# our stand-in for "a city the model has never seen", and it is the only thing
# we are allowed to choose models on.

# %%
blocks_mad = R.spatial_blocks(coords_mad, n_blocks=25)
print("block sizes:", np.bincount(blocks_mad))

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(coords_mad[:, 0], coords_mad[:, 1], c=blocks_mad, s=1, cmap="tab20")
ax.set(title="Madrid split into spatial blocks", xlabel="px_key", ylabel="py_key")
plt.tight_layout()
plt.show()


def block_cv(X, y, blocks, fit_fn, n_folds=5, normalise="robustz", seed=0):
    """Hold out whole geographic blocks, score macro F1 on each."""
    uniq = np.unique(blocks)
    rng = np.random.default_rng(seed)
    order = rng.permutation(uniq)
    folds = np.array_split(order, n_folds)
    scores, oof = [], np.zeros((len(y), 4))
    for f in folds:
        te = np.isin(blocks, f)
        tr = ~te
        nrm = G.CityNormaliser(normalise).fit(X[tr]) if normalise else None
        Xtr = nrm.transform(X[tr]) if nrm else X[tr]
        # The held-out block plays the part of a new city: it gets normalised
        # with its OWN statistics, exactly as Amsterdam will be.
        nte = G.CityNormaliser(normalise).fit(X[te]) if normalise else None
        Xte = nte.transform(X[te]) if nte else X[te]
        m = fit_fn(Xtr, y[tr])
        p = m.predict_proba(Xte)
        oof[te] = p
        scores.append(E.macro_f1(y[te], E.CLASSES[np.argmax(p, 1)]))
    return np.array(scores), oof

# %% [markdown]
# ## 3.2 Which features betray the city?
#
# Train a classifier to separate Madrid pixels from Amsterdam pixels. If it
# succeeds easily, the features carry a strong city signature. The most
# important features in *that* task are the ones most likely to mislead the age
# model, and are candidates for removal.

# %%
auc, imp = G.adversarial_validation(madrid[ALL_FEATS], amsterdam[ALL_FEATS], ALL_FEATS)
print(f"city-vs-city AUC: {auc:.4f}   (1.0 = the two cities are trivially told apart)")
print(imp.head(20).to_string(index=False))

# %% [markdown]
# The joint model tells us the cities are separable; it does not tell us which
# feature is to blame, because importance is shared among correlated features.
# A single-feature AUC does, and the before/after table is the most convincing
# figure in the write-up: 0.5 means a feature is domain-invariant, 1.0 means it
# is a pure city marker and any Madrid split on it is meaningless in Amsterdam.

# %%
Xm_all = madrid[ALL_FEATS].to_numpy(float)
Xa_all = amsterdam[ALL_FEATS].to_numpy(float)
sf_raw = G.single_feature_domain_auc(Xm_all, Xa_all, ALL_FEATS)

# The rankable mask is a STRUCTURAL decision about the features, so it is
# computed once on Madrid - the design city - and shared. Fitting it per city
# lets a column be ranked in one and z-scored in the other, which puts them on
# different scales and reports a domain AUC of 1.000 for a feature that is
# actually domain-invariant. `n_observed_years` did exactly that.
RANKABLE = G.rankable_mask(Xm_all)
print(f"rank-transformable on Madrid: {RANKABLE.sum()} of {len(ALL_FEATS)}; "
      f"too heavily tied to rank: {[ALL_FEATS[j] for j in np.flatnonzero(~RANKABLE)]}")
nm_ = G.CityNormaliser("rank", rankable=RANKABLE).fit(Xm_all).transform(Xm_all)
na_ = G.CityNormaliser("rank", rankable=RANKABLE).fit(Xa_all).transform(Xa_all)
sf_rank = G.single_feature_domain_auc(nm_, na_, ALL_FEATS)
screen = (sf_raw.rename(columns={"domain_auc": "auc_raw"})
          .merge(sf_rank.rename(columns={"domain_auc": "auc_after_rank"}), on="feature")
          .sort_values("auc_raw", ascending=False))
print("worst offenders, raw vs after per-city rank normalisation:")
print(screen.head(25).round(3).to_string(index=False))
print(f"\nmedian domain AUC: raw {screen.auc_raw.median():.3f} -> "
      f"after rank {screen.auc_after_rank.median():.3f}")
print(f"features still above 0.90 after normalisation: "
      f"{(screen.auc_after_rank > 0.90).sum()} of {len(screen)}")
screen.to_csv(C.REPORT_DIR / "domain_auc_screen.csv", index=False)

fig, ax = plt.subplots(figsize=(7, 5))
ax.scatter(screen.auc_raw, screen.auc_after_rank, s=8, alpha=0.5)
ax.plot([0.5, 1], [0.5, 1], "k--", lw=1)
ax.set(xlabel="domain AUC, raw", ylabel="domain AUC, after per-city rank",
       title="Per-city ranking removes the city signature (points below the line)")
plt.tight_layout()
plt.show()

# %%
# Repeat after per-city normalisation: a good normaliser should make the cities
# harder to tell apart without destroying the age signal.
nm = G.CityNormaliser("robustz").fit(Xm_all).transform(Xm_all)
na = G.CityNormaliser("robustz").fit(Xa_all).transform(Xa_all)
auc_z, imp_z = G.adversarial_validation(nm, na, ALL_FEATS)
auc_r, imp_r = G.adversarial_validation(nm_, na_, ALL_FEATS)   # shared mask
print(f"AUC raw {auc:.4f} | robust-z {auc_z:.4f} | rank {auc_r:.4f}")
print("\nstill city-specific after rank normalisation:")
print(imp_r.head(12).to_string(index=False))

# %% [markdown]
# Read the two screens together, because they say different things and the
# difference is the point.
#
# Per-city **rank** normalisation forces every feature's marginal distribution
# to be uniform in both cities, so the *single-feature* domain AUC becomes
# exactly 0.5 by construction. That is not evidence that the domain gap is
# gone. If the *joint* adversarial AUC above is still near 1.0 after ranking,
# the city signature has simply moved into the **correlations between**
# features, where a marginal transform cannot reach it.
#
# That is precisely what CORAL addresses - it matches second-order statistics
# rather than marginals. So the honest conclusion from these two cells is:
# ranking is necessary and cheap, and it is not sufficient. Whether CORAL on
# top actually helps is an empirical question answered in section 3.5, not a
# foregone conclusion: in synthetic testing CORAL slightly *hurt* when the
# shift was close to diagonal.

# %% [markdown]
# ### The drop rule, and why it must be conservative
#
# The tempting rule is "drop the top features by adversarial importance". It is
# wrong. With 59 features the average importance share is 1/59 = 0.017, so any
# threshold near that drops whatever happens to be slightly above average -
# and once normalisation has pushed the joint AUC down to ~0.5, those
# importances are *noise*. In an earlier run this rule dropped three features
# whose single-feature domain AUC was exactly 0.500, i.e. perfectly
# domain-invariant ones.
#
# The principled rule needs two conditions, both of them absolute rather than
# relative: a feature is dropped only if it is **still** strongly
# city-discriminative on its own after normalisation, **and** it is not
# carrying much of the age signal on Madrid. Usually that drops nothing, which
# is the correct answer.

# %%
AUC_CUT = 0.90          # still a city marker after normalisation
IMP_CUT = 0.005         # and contributes little to the Madrid model

X_all = madrid[ALL_FEATS].to_numpy(float)
probe = G.fit_lgbm(G.CityNormaliser("rank").fit_transform(X_all), y_mad,
                   params=dict(n_estimators=300))
age_imp = pd.DataFrame({
    "feature": ALL_FEATS,
    "age_importance": probe.feature_importances_ / max(probe.feature_importances_.sum(), 1)
})
drop_tbl = screen.merge(age_imp, on="feature")
leaky = drop_tbl[(drop_tbl.auc_after_rank > AUC_CUT)
                 & (drop_tbl.age_importance < IMP_CUT)].feature.tolist()
FEATURE_SETS["pruned"] = [c for c in ALL_FEATS if c not in leaky]
print(f"features still above AUC {AUC_CUT} after normalisation: "
      f"{(drop_tbl.auc_after_rank > AUC_CUT).sum()}")
print(f"dropping {len(leaky)}: {leaky if leaky else 'nothing - as expected'}")
print(f"pruned set: {len(FEATURE_SETS['pruned'])} features")
print("\nmost city-specific features that we KEPT because they carry age signal:")
print(drop_tbl.sort_values("auc_after_rank", ascending=False)
      .head(8).round(4).to_string(index=False))

# %% [markdown]
# ## 3.3 Model and feature-set comparison, on Madrid only
#
# Every row of this table is decided with Madrid block-CV. Amsterdam is not
# consulted. This is the table that justifies the final model choice.

# %%
# The comparison runs 4 models x 3 feature sets x 5 blocked folds = 60 fits.
# At full strength (1200 LightGBM rounds, 2000 CatBoost iterations) that is
# tens of minutes on 76k rows, and it buys nothing: we are reading the
# ORDERING of configurations here, not their final scores. So compare cheaply,
# then refit the winner at full strength in section 3.7.
CHEAP_LGB = dict(n_estimators=400, learning_rate=0.06)
CHEAP_CAT = dict(iterations=500, learning_rate=0.06)


def lgbm_balanced(X, y):
    return G.fit_lgbm(X, y, params=CHEAP_LGB,
                      class_weight=dict(zip(range(4), G.balanced_class_weights(y))))

def lgbm_plain(X, y):
    return G.fit_lgbm(X, y, params=CHEAP_LGB)

def cat_balanced(X, y):
    return G.fit_catboost(X, y, params=CHEAP_CAT,
                          class_weights=G.balanced_class_weights(y))

def ordinal_lgbm(X, y):
    return G.OrdinalGBDT(params=CHEAP_LGB).fit(X, y)

import time

FEATURE_SETS_TO_TRY = [k for k in ("official60", "everything", "pruned")
                       if len(FEATURE_SETS.get(k, []))]
# Drop a duplicate: with the pkl-only fallback, "everything" == "official60".
seen, uniq = set(), []
for k in FEATURE_SETS_TO_TRY:
    sig = tuple(FEATURE_SETS[k])
    if sig not in seen:
        seen.add(sig)
        uniq.append(k)
    else:
        print(f"  (skipping '{k}' - identical to an earlier set)")
FEATURE_SETS_TO_TRY = uniq

rows = []
for fs_name in FEATURE_SETS_TO_TRY:
    X = madrid[FEATURE_SETS[fs_name]].to_numpy(float)
    models = [("lgbm", lgbm_balanced), ("lgbm_unweighted", lgbm_plain),
              ("lgbm_ordinal", ordinal_lgbm)]
    if fs_name == FEATURE_SETS_TO_TRY[0]:
        # CatBoost costs 35 min per feature set against LightGBM's 10 and came
        # in about 0.008 behind on every set, so it runs on one set only. It
        # still ships as the second required model.
        models.append(("catboost", cat_balanced))
    for mdl_name, fn in models:
        t0 = time.time()
        s, oof = block_cv(X, y_mad, blocks_mad, fn)
        rows.append(dict(features=fs_name, model=mdl_name, n_feat=X.shape[1],
                         macro_f1=s.mean(), std=s.std(ddof=1),
                         minutes=(time.time() - t0) / 60))
        print(f"  {fs_name:<12} {mdl_name:<16} {s.mean():.4f} +- "
              f"{s.std(ddof=1):.4f}   ({rows[-1]['minutes']:.1f} min)")
table = pd.DataFrame(rows).sort_values("macro_f1", ascending=False)
print("\n" + table.round(4).to_string(index=False))

# %% [markdown]
# ### Does per-city normalisation actually help? Decide it on Madrid.
#
# This has to be settled here and not on the Amsterdam score, because the
# Amsterdam score is not ours to choose on. `block_cv` is the right instrument:
# it normalises the held-out block with the block's OWN statistics, exactly as
# Amsterdam will be normalised, so each fold is a miniature transfer problem.
#
# Worth doing carefully, because the intuition ("absolute reflectance does not
# transfer, so normalise it away") is strong and may still be wrong. Per-city
# normalisation also destroys genuine between-city differences in LEVEL, and if
# a city's overall darkness happens to correlate with its class mix, removing
# the level removes a real cue along with the nuisance.

# %%
norm_rows = []
X_best = madrid[FEATURE_SETS[table.iloc[0].features]].to_numpy(float)
for norm in (None, "robustz", "rank"):
    s, _ = block_cv(X_best, y_mad, blocks_mad, lgbm_plain, normalise=norm)
    norm_rows.append(dict(normalise=str(norm), macro_f1=s.mean(),
                          std=s.std(ddof=1)))
    print(f"  normalise={str(norm):8} {s.mean():.4f} +- {s.std(ddof=1):.4f}")
norm_table = pd.DataFrame(norm_rows).sort_values("macro_f1", ascending=False)
print("\n" + norm_table.round(4).to_string(index=False))
BEST_NORM = None if norm_table.iloc[0].normalise == "None" else norm_table.iloc[0].normalise
print(f"\nnormalisation chosen on Madrid: {BEST_NORM}")
print("Whatever Amsterdam says later, this is the choice the rules allow us to")
print("make. If Amsterdam disagrees, that is a finding to report, not a reason")
print("to switch.")

# %% [markdown]
# ## 3.4 The decision rule matters as much as the model
#
# Macro F1 is not maximised by `argmax`. Three cheap post-processing steps,
# each measured separately on Madrid block-CV:
#
# * **class weights** tuned by coordinate ascent on out-of-fold probabilities;
# * **prior correction** towards the class prior EM estimates from *unlabelled*
#   target features;
# * **spatial smoothing** of the probabilities over neighbouring pixels.

# %%
BEST_FS = table.iloc[0].features
X_mad = madrid[FEATURE_SETS[BEST_FS]].to_numpy(float)
X_ams = amsterdam[FEATURE_SETS[BEST_FS]].to_numpy(float)
scores, oof = block_cv(X_mad, y_mad, blocks_mad, lgbm_balanced)
print(f"baseline argmax, Madrid block-CV: {E.macro_f1(y_mad, E.CLASSES[np.argmax(oof, 1)]):.4f}")

w = E.tune_class_weights(oof, y_mad)
print("tuned class weights:", np.round(w, 3))
print("after weighting:   ",
      f"{E.macro_f1(y_mad, E.predict_weighted(oof, w)):.4f}")

for radius in (2.0, 4.0, 8.0):
    sm = E.spatial_smooth(oof, coords_mad, radius=radius, alpha=0.5)
    print(f"spatial smoothing r={radius:>4}: "
          f"{E.macro_f1(y_mad, E.predict_weighted(sm, w)):.4f}")

# %%
print(pd.DataFrame(E.confusion_proportions(y_mad, E.predict_weighted(oof, w)),
                   index=[f"true {c}" for c in E.CLASSES],
                   columns=[f"pred {c}" for c in E.CLASSES]).round(3).to_string())

# %% [markdown]
# ## 3.5 The simulated-shift harness - choosing the decision rule on Madrid
#
# This is the most important cell in the notebook, and the thing I would most
# want the team to copy.
#
# None of the adaptation tricks above can be chosen honestly without a way to
# *simulate* Amsterdam from Madrid. So: hold out a spatial block, and build four
# versions of it.
#
# * `clean` - as-is;
# * `label` - resampled to a random Dirichlet class prior (pure label shift);
# * `covariate` - every feature quantile-mapped onto Amsterdam's own marginal
#   (pure covariate shift, using `X_amsterdam` only, no labels);
# * `both`.
#
# Then score every candidate decision rule on all four. A rule that helps under
# label shift can lose badly under covariate shift, and without this there is no
# way to tell which regime you are actually in. **Ship only rules that are
# non-negative in the worst case.**
#
# In synthetic testing this harness caught a trap: free SLD-EM prior estimation
# looked excellent under label shift and lost 0.19 macro F1 on average (0.35 at
# worst) under covariate shift, because EM assumes label shift only. The
# macro-F1 plug-in rule driven by the *source* prior was the robust winner.

# %%
# Note the UNWEIGHTED model here. The whole point of this cell is to choose the
# operating point, so the probabilities going into it must be an honest estimate
# of the Madrid posterior. class_weight='balanced' bakes a decision into the fit,
# where no decision rule can undo it.
harness_df, harness_table = E.shift_harness(
    lgbm_plain, X_mad, y_mad, blocks_mad, X_ams,
    E.DEFAULT_STRATEGIES, n_folds=5, n_rep=4, seed=C.RANDOM_SEED)
harness_table.to_csv(C.REPORT_DIR / "shift_harness.csv")

# %%
# The criterion is the PER-REGIME worst case: for each rule, its deficit
# against argmax in whichever regime it handles worst. An earlier version used
# a column that compared each rule's worst cell anywhere against argmax's worst
# cell anywhere - two different regimes and two different folds - so a rule
# could look safe merely because its bad regime was also argmax's bad regime.
WORST = "delta_vs_argmax_worst_regime"
safe = harness_table[harness_table.get(WORST, 0) >= 0]
print("rules that never lose to plain argmax in ANY regime:")
print(safe.round(4).to_string() if len(safe) else "  none - keep plain argmax")
print("\nrules to avoid:")
risky = harness_table[harness_table.get(WORST, 0) < -0.01]
print(risky.round(4).to_string() if len(risky) else "  none")

CHOSEN_RULE = (harness_table[WORST].idxmax()
               if WORST in harness_table else "argmax")
print(f"\nrule chosen on Madrid alone: {CHOSEN_RULE}")
DECISION = E.DEFAULT_STRATEGIES[CHOSEN_RULE]

# %% [markdown]
# ## 3.6 The official few-shot protocol
#
# `y_amsterdam` enters here and only here: to draw the n labelled pixels per
# class that the rules allow, and to score the remainder. Nothing above was
# tuned on it.

# %%
N_LIST = (5, 25, 50, 100, 200)
N_REPEATS = 12
# Score a fixed random 8,000-pixel sample of Amsterdam rather than all 25,992.
# Full scoring cost 19 s per draw - 32 minutes for one variant and hours for the
# table - and bought nothing: the variants that ignore the labelled sample have
# a standard deviation of 0.0001 across draws. The subsample is fixed, so every
# row of the table is compared on identical pixels. Re-run the winner on all
# 25,992 once, at the end, for the number that goes in the report.
EVAL_SUBSAMPLE = 8000
src_prior = np.array([(y_mad == c).mean() for c in E.CLASSES])
print('Madrid class prior:', np.round(src_prior, 3))

# Cost control. Most variants cache the Madrid-trained model across repeats, so
# 20 draws are cheap. `upweighted shots, retrained` refits from scratch on every
# draw - 100 full fits on 76k rows - so it gets cheap params and fewer draws,
# and its std is correspondingly noisier. That is noted in the table.
REPEATS = {"upweighted shots, retrained": 5}

variants = {
    # the two baselines everything must beat
    "organisers' prototype (Euclidean, Madrid scaler)": G.mc_prototype_official(),
    "pure retrieval k-NN (no source model)": G.mc_knn(),
    # zero-shot: does normalisation alone carry us?
    "zero-shot, no normalisation": G.mc_zero_shot(normalise=None),
    "zero-shot, robust z": G.mc_zero_shot(normalise="robustz"),
    "zero-shot, rank": G.mc_zero_shot(normalise="rank"),
    "zero-shot, robust z + CORAL": G.mc_zero_shot(normalise="robustz", use_coral=True),
    # the prototype metric ablation. On synthetic anisotropic data the pooled
    # WITHIN-class Mahalanobis metric beat the organisers' Euclidean prototype
    # by +0.11 macro F1 at n=25 - the largest single gain of anything tried.
    # The total-covariance version is worse because the target city's total
    # covariance is dominated by BETWEEN-class spread, so it shrinks exactly the
    # directions that separate the classes. Confirm the ordering here on the
    # real features before believing it.
    "prototype, robust z + euclidean": G.mc_prototype(metric="euclidean"),
    "prototype, rank + euclidean": G.mc_prototype(normalise="rank", metric="euclidean"),
    "prototype, robust z + total-cov Mahalanobis": G.mc_prototype(metric="total"),
    "prototype, robust z + within-class Mahalanobis": G.mc_prototype(metric="within"),
    # using the labelled sample with the source model
    "recalibrated (stratification-corrected)": G.mc_recalibrated(normalise="robustz"),
    "refit leaf values (structure frozen)": G.mc_refit(),
    "upweighted shots, retrained": G.mc_shot_weighted(
        normalise="robustz", fit=lambda X, y, **kw: G.fit_lgbm(X, y, params=CHEAP_LGB, **kw)),
    "catboost recalibrated": G.mc_recalibrated(fit=G.fit_catboost, normalise="robustz"),
}

protocol = {}
for name, mc in variants.items():
    print(f"\n{name}")
    reps = REPEATS.get(name, N_REPEATS)
    t0 = time.time()
    protocol[name] = E.evaluate_protocol(
        mc, X_mad, y_mad, X_ams, y_ams, n_list=N_LIST, n_repeats=reps,
        decision=lambda p: DECISION(p, src_prior),
        eval_subsample=EVAL_SUBSAMPLE)
    print(f"  ({reps} draws per n, {(time.time() - t0) / 60:.1f} min)")

# %%
summary = pd.DataFrame({name: {n: r["mean"] for n, r in res.items()}
                        for name, res in protocol.items()}).T
summary["average"] = summary.mean(axis=1)
summary["peak"] = summary[list(N_LIST)].max(axis=1)
print(summary.round(4).sort_values("average", ascending=False).to_string())

# %%
fig, ax = plt.subplots(figsize=(8, 5))
for name, res in protocol.items():
    m = [res[n]["mean"] for n in N_LIST]
    s = [res[n]["std"] for n in N_LIST]
    ax.errorbar(np.log2(N_LIST), m, yerr=s, marker="o", capsize=3, label=name)
ax.axhline(0.618, ls="--", c="k", lw=1, label="published Madrid baseline 0.618")
ax.set(xlabel="log2(labelled Amsterdam pixels per class)", ylabel="macro F1 on Amsterdam",
       title="Macro F1 against labelling budget")
ax.legend(fontsize=7)
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 3.7 Save the probabilities for the ensemble
#
# The LLM notebook loads these and blends them with its own predictions.

# %%
import pickle

best_name = summary.sort_values("average", ascending=False).index[0]
print("best classical variant on the protocol:", best_name)

nrm_m = G.CityNormaliser("robustz").fit(X_mad)
nrm_a = G.CityNormaliser("robustz").fit(X_ams)
final = G.fit_lgbm(nrm_m.transform(X_mad), y_mad,
                   class_weight=dict(zip(range(4), G.balanced_class_weights(y_mad))))
final_cat = G.fit_catboost(nrm_m.transform(X_mad), y_mad,
                           class_weights=G.balanced_class_weights(y_mad))

out = {
    "feature_set": BEST_FS,
    "feature_names": FEATURE_SETS[BEST_FS],
    "madrid_oof_proba": oof,
    "madrid_blocks": blocks_mad,
    "class_weights": w,
    "amsterdam_proba_lgbm": final.predict_proba(nrm_a.transform(X_ams)),
    "amsterdam_proba_catboost": final_cat.predict_proba(nrm_a.transform(X_ams)),
    "protocol": {k: {n: dict(mean=r["mean"], std=r["std"],
                             per_class=r["per_class"].tolist())
                     for n, r in v.items()} for k, v in protocol.items()},
    # Plain lists, not a DataFrame: this file is read on Colab, whose pandas
    # 2.x cannot unpickle a pandas 3 DataFrame with string columns
    # ("StringDtype.__init__() takes from 1 to 2 positional arguments").
    # Rebuild with pd.DataFrame(out["madrid_block_cv_table"]).
    "madrid_block_cv_table": {c: table[c].tolist() for c in table.columns},
    "adversarial_auc": dict(raw=auc, robustz=auc_z, rank=auc_r),
    "dropped_city_specific": leaky,
}
with open(C.PROC_DIR / "gbdt_outputs.pkl", "wb") as fh:
    pickle.dump(out, fh)
final.booster_.save_model(str(C.PROC_DIR / "model_lgbm.txt"))
final_cat.save_model(str(C.PROC_DIR / "model_catboost.cbm"))
print("saved:", C.PROC_DIR / "gbdt_outputs.pkl")