"""
The two winning models from our Hack4Dev building-age entry, self-contained.

Needs only numpy, pandas, pyarrow, scikit-learn and lightgbm - none of the
project's src/ code. Every function here is a line-for-line port of the code
that produced the reported numbers (src/gbdt.py, src/evaluation.py); the
export script checked that the ports give identical results.

  ZERO-SHOT     LightGBM trained on Madrid only. Never sees an Amsterdam label.
                Load it from models/zero_shot_lgbm.txt, or retrain it.
  RETRAINED     A fresh LightGBM on all of Madrid plus the n labelled Amsterdam
                pixels per class that the protocol hands you, those pixels
                weighted 30x. Refit for every draw of labelled pixels, so there
                is no single saved model - models/retrained_*.txt are two
                example draws, saved so you can inspect one.

Both use the same decision rule: tau = 0.5 logit adjustment by the Madrid class
prior, then the macro-F1 plug-in rule. Chosen on Madrid alone.

Command line (run from this folder):

    python winners.py zero-shot                  # score the saved model
    python winners.py retrained --n 25 --rep 0   # one draw, as reported
    python winners.py protocol --model retrained --reps 5
    python winners.py protocol --model zero-shot --reps 3

Python:

    import winners as W
    d = W.load_data()
    model = W.load_zero_shot()
    labels = W.decide(W.predict(model, d["X_ams"]), d["madrid_prior"])
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd

HERE = Path(__file__).resolve().parent
DATA = HERE / "data" / "processed"
MODELS = HERE / "models"

CLASSES = np.array([1, 2, 3, 4])   # 1 = oldest ... 4 = newest; see README
SEED = 42

# Full-strength LightGBM (src/gbdt.py LGBM_PARAMS) - used by the zero-shot model.
LGBM_PARAMS = dict(
    objective="multiclass", num_class=4, learning_rate=0.03, n_estimators=1200,
    num_leaves=31, max_depth=-1, min_child_samples=80, subsample=0.8,
    subsample_freq=1, colsample_bytree=0.6, reg_alpha=0.5, reg_lambda=5.0,
    max_bin=127, n_jobs=-1, random_state=SEED, verbosity=-1,
)
# The retrained variant was measured with notebook 3's cheaper settings,
# because it refits on every draw. Same params otherwise.
RETRAIN_OVERRIDES = dict(n_estimators=400, learning_rate=0.06)
SHOT_WEIGHT = 30.0


# --------------------------------------------------------------------------- #
# Data
# --------------------------------------------------------------------------- #
def feature_names() -> list[str]:
    return json.loads((MODELS / "meta.json").read_text())["feature_names"]


def load_data(data_dir: Path | str = DATA) -> dict:
    """Madrid and Amsterdam feature matrices in the exact training column order."""
    data_dir = Path(data_dir)
    feats = feature_names()
    mad = pd.read_parquet(data_dir / "madrid_features.parquet")
    ams = pd.read_parquet(data_dir / "amsterdam_features.parquet")
    y_mad = mad["age_class"].to_numpy()
    return dict(
        X_mad=mad[feats].to_numpy(float), y_mad=y_mad,
        X_ams=ams[feats].to_numpy(float), y_ams=ams["age_class"].to_numpy(),
        coords_ams=ams[["px_key", "py_key"]].to_numpy(float),
        madrid_prior=np.array([(y_mad == c).mean() for c in CLASSES]),
        feature_names=feats,
    )


# --------------------------------------------------------------------------- #
# Decision rule: tau-logit adjustment, then the macro-F1 plug-in
# --------------------------------------------------------------------------- #
def logit_adjust(proba, train_prior, tau=0.5):
    p = np.asarray(proba, dtype=float)
    adj = p / np.clip(np.asarray(train_prior, float), 1e-12, None) ** float(tau)
    return adj / np.clip(adj.sum(axis=1, keepdims=True), 1e-12, None)


def macro_f1_plugin(proba, prior, n_iter=50):
    """Ramaswamy et al.'s plug-in rule for macro F1. Uses no target labels, but
    it does look at the predicted class shares of the batch it is given, so
    call it on the whole set you want labelled, not row by row."""
    P = np.asarray(proba, dtype=float)
    pi = np.asarray(prior, dtype=float)
    pi = pi / pi.sum()
    f1 = np.zeros(4)
    best = None
    for it in range(n_iter):
        q = np.clip(np.bincount(np.argmax(P, axis=1), minlength=4) / len(P), 1e-6, None) \
            if it == 0 else q_new
        a = 2.0 / np.clip(pi + q, 1e-9, None)
        lab_idx = np.argmax(a[None, :] * (P - f1[None, :] / 2.0), axis=1)
        q_new = np.clip(np.bincount(lab_idx, minlength=4) / len(P), 1e-6, None)
        tp = np.array([P[lab_idx == k, k].sum() for k in range(4)]) / len(P)
        f1_new = 2 * tp / np.clip(pi + q_new, 1e-9, None)
        obj = float(np.mean(f1_new))
        if best is None or obj > best[0]:
            best = (obj, CLASSES[lab_idx].copy())
        if np.max(np.abs(f1_new - f1)) < 1e-6 and it > 0:
            break
        f1 = f1_new
    return best[1]


def decide(proba, madrid_prior, tau=0.5):
    """Class labels (1-4) from probabilities. Plain argmax would be
    CLASSES[proba.argmax(1)]; this rule was worth about +0.01 on Madrid."""
    src = np.asarray(madrid_prior, float)
    src = src / src.sum()
    return macro_f1_plugin(logit_adjust(proba, src, tau), src)


def macro_f1(y_true, y_pred) -> float:
    from sklearn.metrics import f1_score
    return float(f1_score(y_true, y_pred, labels=CLASSES, average="macro",
                          zero_division=0))


# --------------------------------------------------------------------------- #
# Models
# --------------------------------------------------------------------------- #
def balanced_class_weights(y):
    y = np.asarray(y)
    counts = np.array([(y == c).sum() for c in CLASSES], dtype=float)
    return len(y) / (4.0 * np.clip(counts, 1, None))


def _fit_lgbm(X, y, params, class_weight=None, sample_weight=None):
    import lightgbm as lgb
    p = dict(params)
    if class_weight is not None:
        p["class_weight"] = class_weight
    m = lgb.LGBMClassifier(**p)
    m.fit(np.asarray(X), np.asarray(y) - 1, sample_weight=sample_weight)
    return m


def train_zero_shot(X_mad, y_mad):
    """Madrid only, class-balanced, raw features (no per-city scaling)."""
    cw = balanced_class_weights(y_mad)
    return _fit_lgbm(X_mad, y_mad, LGBM_PARAMS,
                     class_weight=dict(zip(range(4), cw)))


def robustz_fit(X):
    X = np.asarray(X, dtype=np.float64)
    med = np.median(X, axis=0)
    q75, q25 = np.percentile(X, [75, 25], axis=0)
    iqr = np.where((q75 - q25) > 1e-12, q75 - q25, 1.0)
    return med, iqr


def robustz(X, stats):
    med, iqr = stats
    return (np.asarray(X, dtype=np.float64) - med) / iqr


def train_retrained(X_mad, y_mad, X_ams_all, X_shot, y_shot,
                    shot_weight=SHOT_WEIGHT):
    """
    Retrain on Madrid + the labelled Amsterdam shots, shots upweighted.

    Each city is robust-z scaled with its OWN median/IQR. Amsterdam's are taken
    from its whole unlabelled feature matrix (features only, no labels) - the
    protocol passes that in, and the organisers' baseline does the same.
    Returns (model, amsterdam_stats); predict with predict_retrained().
    """
    mad_stats = robustz_fit(X_mad)
    ams_stats = robustz_fit(X_ams_all)
    X = np.vstack([robustz(X_mad, mad_stats), robustz(X_shot, ams_stats)])
    y = np.r_[np.asarray(y_mad), np.asarray(y_shot)]
    w = np.r_[np.ones(len(y_mad)), np.full(len(y_shot), shot_weight)]
    w = w * balanced_class_weights(y)[y - 1]
    model = _fit_lgbm(X, y, {**LGBM_PARAMS, **RETRAIN_OVERRIDES}, sample_weight=w)
    return model, ams_stats


def predict(model, X):
    """Class probabilities, columns = classes 1..4. Accepts a fitted
    LGBMClassifier or a lightgbm.Booster loaded from models/*.txt."""
    if hasattr(model, "predict_proba"):
        return model.predict_proba(np.asarray(X))
    return model.predict(np.asarray(X))


def predict_retrained(model, ams_stats, X):
    return predict(model, robustz(X, ams_stats))


def load_zero_shot():
    import lightgbm as lgb
    return lgb.Booster(model_file=str(MODELS / "zero_shot_lgbm.txt"))


def load_retrained_example(n: int):
    """One saved retrained model (draw rep 0 at this n) and its Amsterdam stats."""
    import lightgbm as lgb
    booster = lgb.Booster(model_file=str(MODELS / f"retrained_n{n}_rep0_lgbm.txt"))
    z = np.load(MODELS / "amsterdam_robustz_stats.npz")
    return booster, (z["median"], z["iqr"])


# --------------------------------------------------------------------------- #
# The official few-shot protocol
# --------------------------------------------------------------------------- #
def draw_shots(y_tgt, n_per_class, rng):
    """Stratified n-per-class sample - identical to src/evaluation.draw_shots."""
    picks = []
    for c in CLASSES:
        pool = np.flatnonzero(np.asarray(y_tgt) == c)
        if len(pool):
            picks.append(rng.choice(pool, size=min(n_per_class, len(pool)),
                                    replace=False))
    return np.concatenate(picks)


def shots_for(y_ams, n, rep):
    """The exact draw we reported for (n, rep)."""
    return draw_shots(y_ams, n, np.random.default_rng(SEED + 1000 * n + rep))


def score_draw(d, model_name, n, rep, zero_shot_model=None):
    """Macro F1 on all of Amsterdam minus this draw's labelled pixels."""
    shot = shots_for(d["y_ams"], n, rep)
    keep = np.setdiff1d(np.arange(len(d["y_ams"])), shot)
    if model_name == "zero-shot":
        m = zero_shot_model or load_zero_shot()
        P = predict(m, d["X_ams"][keep])
    else:
        m, st = train_retrained(d["X_mad"], d["y_mad"], d["X_ams"],
                                d["X_ams"][shot], d["y_ams"][shot])
        P = predict_retrained(m, st, d["X_ams"][keep])
    return macro_f1(d["y_ams"][keep], decide(P, d["madrid_prior"]))


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("what", choices=["zero-shot", "retrained", "protocol"])
    ap.add_argument("--model", choices=["zero-shot", "retrained"], default="retrained")
    ap.add_argument("--n", type=int, default=25)
    ap.add_argument("--rep", type=int, default=0)
    ap.add_argument("--reps", type=int, default=5)
    a = ap.parse_args()
    d = load_data()
    if a.what == "zero-shot":
        P = predict(load_zero_shot(), d["X_ams"])
        print(f"zero-shot, all {len(P):,} Amsterdam pixels: macro F1 "
              f"{macro_f1(d['y_ams'], decide(P, d['madrid_prior'])):.4f}")
    elif a.what == "retrained":
        print(f"retrained, n={a.n} rep={a.rep}: macro F1 "
              f"{score_draw(d, 'retrained', a.n, a.rep):.4f}")
    else:
        zs = load_zero_shot() if a.model == "zero-shot" else None
        for n in (5, 25, 50, 100, 200):
            s = [score_draw(d, a.model, n, r, zs) for r in range(a.reps)]
            print(f"{a.model} n={n:<4d} macro F1 {np.mean(s):.4f} "
                  f"+- {np.std(s, ddof=1) if len(s) > 1 else 0:.4f}", flush=True)


if __name__ == "__main__":
    main()