# %% [markdown]
# # The two winners, end to end
#
# Building age class (1-4) per 30 m Landsat pixel. Train on Madrid, predict
# Amsterdam, scored by macro F1 at n = 5, 25, 50, 100, 200 labelled Amsterdam
# pixels per class.
#
# * **Zero-shot LightGBM** - trained on Madrid only, never sees an Amsterdam
#   label. Best model if you may not use Amsterdam at all: **0.593**.
# * **Retrained LightGBM** - refit on Madrid plus the n labelled Amsterdam pixels
#   per class the protocol gives you, those pixels weighted 30x. Best average
#   (**0.656**) and best peak (**0.720** at n=200) over the whole protocol.
#
# Everything here runs on a laptop CPU. The heavy lifting is in `winners.py`;
# this notebook just calls it. All numbers are on ALL of Amsterdam minus the
# labelled pixels of each draw.

# %%
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path.cwd() if (Path.cwd() / "winners.py").exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import winners as W

d = W.load_data()
print("Madrid   ", d["X_mad"].shape, " Amsterdam", d["X_ams"].shape,
      " features", len(d["feature_names"]))
print("Madrid class shares   ", np.round(d["madrid_prior"], 3))
print("Amsterdam class shares", np.round([(d["y_ams"] == c).mean() for c in W.CLASSES], 3))

# %% [markdown]
# ## 1. Zero-shot LightGBM: Madrid only
#
# 1,200 boosting rounds, class-balanced, raw features (per-city scaling made
# things *worse* on Madrid's own spatial cross-validation, so none is applied).
# The final step is the decision rule: tau = 0.5 logit adjustment by the Madrid
# class prior, then the macro-F1 plug-in rule. It uses no Amsterdam labels, only
# the predicted class shares of the batch being labelled.

# %%
zs = W.load_zero_shot()
P = W.predict(zs, d["X_ams"])
lab = W.decide(P, d["madrid_prior"])
from sklearn.metrics import confusion_matrix, f1_score
print(f"macro F1, all {len(lab):,} Amsterdam pixels: {W.macro_f1(d['y_ams'], lab):.4f}")
print(f"  same model, plain argmax instead of the rule: "
      f"{W.macro_f1(d['y_ams'], W.CLASSES[P.argmax(1)]):.4f}")
print("per-class F1:", np.round(f1_score(d["y_ams"], lab, labels=W.CLASSES, average=None), 3))
cm = confusion_matrix(d["y_ams"], lab, labels=W.CLASSES, normalize="true")
print(pd.DataFrame(cm, index=[f"true {c}" for c in W.CLASSES],
                   columns=[f"pred {c}" for c in W.CLASSES]).round(3).to_string())

# %% [markdown]
# ## 2. Retrained LightGBM: one draw
#
# The protocol hands you n labelled Amsterdam pixels per class. Here n = 25,
# draw 0 - the exact pixels behind the reported number. A fresh LightGBM
# (400 rounds) is fit on all 76k Madrid pixels plus these 100, the 100 weighted
# 30x. Each city is robust-z scaled with its own median and IQR, Amsterdam's
# taken from its unlabelled features. The 100 pixels are removed before scoring.

# %%
import time
n, rep = 25, 0
shot = W.shots_for(d["y_ams"], n, rep)
keep = np.setdiff1d(np.arange(len(d["y_ams"])), shot)
print(f"{len(shot)} labelled pixels, classes {np.bincount(d['y_ams'][shot])[1:]}; "
      f"scoring the other {len(keep):,}")
t0 = time.time()
model, ams_stats = W.train_retrained(d["X_mad"], d["y_mad"], d["X_ams"],
                                     d["X_ams"][shot], d["y_ams"][shot])
P = W.predict_retrained(model, ams_stats, d["X_ams"][keep])
lab = W.decide(P, d["madrid_prior"])
print(f"trained in {time.time() - t0:.0f} s; macro F1 {W.macro_f1(d['y_ams'][keep], lab):.4f}")
print("per-class F1:", np.round(f1_score(d["y_ams"][keep], lab, labels=W.CLASSES, average=None), 3))

# %% [markdown]
# ## 3. The full protocol, as reported
#
# Every n, several draws, scored on all of Amsterdam. These come from
# `results/full_score_*.json` (produced by `tools/full_score.py`); to recompute
# them run `python winners.py protocol --model retrained --reps 5` (~35 min).
# The zero-shot and organisers' rows come from an earlier full-city run
# (`results/logs/03_gbdt_models_fullscore_run.log`, 20 draws per n).

# %%
N = [5, 25, 50, 100, 200]
R = ROOT / "results"
rows = {}
for key, label in [("upweighted", "retrained LightGBM (5 draws per n)"),
                   ("prototype", "prototype, within-class Mahalanobis (12 draws)")]:
    s = json.loads((R / f"full_score_{key}.json").read_text())["summary"]["full"]
    rows[label] = {n: s[str(n)]["mean"] for n in N} | {"std@25": s["25"]["std"]}
rows["zero-shot LightGBM (20 draws)"] = dict(zip(N, [0.5928, 0.5925, 0.5922, 0.5918, 0.5903]))
rows["organisers' prototype, our re-run on 393 features (20 draws)"] = dict(zip(N, [0.5066, 0.5821, 0.5965, 0.6034, 0.5972]))
table = pd.DataFrame(rows).T
table["average"] = table[N].mean(axis=1)
table["peak"] = table[N].max(axis=1)
print(table.round(4).to_string())

# %%
import matplotlib.pyplot as plt

INK, INK2, MUTED, GRID, SURF = "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#fcfcfb"
series = [  # fixed categorical order: blue, orange, aqua
    ("upweighted", "Retrained LightGBM", "#2a78d6"),
    ("prototype", "Prototype (within-class Mahalanobis)", "#eb6834"),
]
fig, ax = plt.subplots(figsize=(8.5, 5), facecolor=SURF)
ax.set_facecolor(SURF)
x = np.log2(N)
for key, name, col in series:
    s = json.loads((R / f"full_score_{key}.json").read_text())["summary"]["full"]
    m = np.array([s[str(n)]["mean"] for n in N])
    e = np.array([s[str(n)]["std"] for n in N])
    ax.fill_between(x, m - e, m + e, color=col, alpha=0.12, lw=0)
    ax.plot(x, m, color=col, lw=2, marker="o", ms=8, mec=SURF, mew=2, label=name)
    dy = 7 if key == "upweighted" else -7      # keep the two end labels apart
    ax.annotate(name, (x[-1], m[-1]), xytext=(10, dy), textcoords="offset points",
                va="center", color=INK2, fontsize=9)
zs_m = np.array(list(rows["zero-shot LightGBM (20 draws)"].values()))
ax.plot(x, zs_m, color="#1baf7a", lw=2, marker="o", ms=8, mec=SURF, mew=2,
        label="Zero-shot LightGBM (Madrid only)")
ax.annotate("Zero-shot LightGBM", (x[-1], zs_m[-1]), xytext=(10, -9),
            textcoords="offset points", va="center", color=INK2, fontsize=9)
base = np.array(list(rows["organisers' prototype, our re-run on 393 features (20 draws)"].values()))
ax.plot(x, base, color=MUTED, lw=1.5, ls="--", label="Organisers' prototype (our re-run)")
ax.annotate("Organisers' prototype (re-run)", (x[-1], base[-1]), xytext=(10, 8),
            textcoords="offset points", va="center", color=MUTED, fontsize=9)
ax.set_xticks(x, [str(n) for n in N])
ax.set_xlabel("Labelled Amsterdam pixels per class (log scale)", color=INK2)
ax.set_ylabel("Macro F1, all of Amsterdam", color=INK2)
ax.set_title("Macro F1 against labelling budget", loc="left", color=INK, fontsize=12)
ax.grid(axis="y", color=GRID, lw=0.8)
ax.set_axisbelow(True)
for sp in ("top", "right"):
    ax.spines[sp].set_visible(False)
for sp in ("left", "bottom"):
    ax.spines[sp].set_color(GRID)
ax.tick_params(colors=MUTED)
ax.set_xlim(x[0] - 0.3, x[-1] + 2.6)
ax.legend(loc="lower right", frameon=False, fontsize=8, labelcolor=INK2)
plt.tight_layout()
fig.savefig(R / "f1_vs_n_full_city.png", dpi=160, facecolor=SURF)
plt.show()

# %% [markdown]
# ## 4. Using them on your own data
#
# Your features must be the same 393 columns, in `models/meta.json` order
# (`W.feature_names()`); `src/features.py` / notebook 1 builds them from the
# raw hackathon parquet files.
#
# ```python
# # zero-shot: no target labels at all
# P = W.predict(W.load_zero_shot(), X_target)
# labels = W.decide(P, madrid_prior)            # or P.argmax(1) + 1
#
# # retrained: you have n labelled target pixels per class
# model, stats = W.train_retrained(X_mad, y_mad, X_target_all, X_shot, y_shot)
# labels = W.decide(W.predict_retrained(model, stats, X_rest), madrid_prior)
# ```
#
# To compare with your own model on identical draws, use `W.shots_for(y_ams, n,
# rep)` for the labelled pixels and score on the rest of Amsterdam - that is the
# setup behind every number in the table above.