"""
Model 1 and Model 2 of the submission: LightGBM and CatBoost, plus the
cross-city plumbing they share with the LLM track.

Contents
  * per-city normalisation (rank / robust z)        - unsupervised, no labels
  * CORAL covariance alignment                      - unsupervised, no labels
  * adversarial validation                          - finds city-specific features
  * macro-F1 eval metrics for both libraries
  * ordinal (cumulative-binary) decomposition       - the classes are ordered
  * several `modelCreate` adapters for the few-shot protocol
"""
from __future__ import annotations

import numpy as np
from sklearn.metrics import f1_score

CLASSES = np.array([1, 2, 3, 4])
SEED = 42


# --------------------------------------------------------------------------- #
# Per-city normalisation.  Fit one of these per city on its own features; it
# needs no labels, so it is legal for Amsterdam.
# --------------------------------------------------------------------------- #
def rankable_mask(X, min_unique: int = 10, max_tie_share: float = 0.5):
    """
    Which columns are safe to rank-transform, judged on the DESIGN city only.

    A column is unsafe when one value dominates it. The test that matters is
    the size of the biggest tie, not the number of distinct values:
    `first_observed_year` has 20 distinct values but is 1984 for 98.4% of
    pixels, and ranking it turned a domain AUC of 0.507 into 0.997.
    """
    X = np.asarray(X, dtype=np.float64)
    out = np.empty(X.shape[1], dtype=bool)
    for j in range(X.shape[1]):
        col = X[:, j]
        vals, counts = np.unique(col, return_counts=True)
        out[j] = (len(vals) >= min_unique
                  and counts.max() / len(col) < max_tie_share)
    return out


class CityNormaliser:
    """
    mode="rank"    : map each feature to its within-city percentile (0-1).
                     Kills every city-level offset and gain, including ones no
                     linear method would catch. Cost: it also erases genuine
                     between-city differences in level, and it is distorted
                     when the two cities have different class proportions.
    mode="robustz" : (x - median) / IQR per city. Gentler, keeps shape.
    mode=None      : identity. The right answer on this data (see
                     TREE_NORMALISE), so it has to be a real, explicit mode.

    `mode=None` used to fall through `if self.mode != "rank"` in `transform`
    and apply robust z anyway. Every caller that asked for no normalisation got
    robust z instead, silently - `mc_refit`, whose default is `TREE_NORMALISE`
    (None), was measured that way. An unrecognised mode now raises rather than
    quietly becoming one of the other two.
    """

    MODES = (None, "none", "robustz", "rank")

    def __init__(self, mode: str | None = "robustz", n_grid: int = 1001,
                 min_unique: int = 10, max_tie_share: float = 0.5,
                 rankable=None):
        if mode not in self.MODES:
            raise ValueError(f"unknown normalise mode {mode!r}; "
                             f"expected one of {self.MODES}")
        self.mode = None if mode == "none" else mode
        self.n_grid = n_grid
        self.min_unique = min_unique
        self.max_tie_share = max_tie_share
        # Which columns may be rank-transformed. This is a STRUCTURAL decision
        # about the features and it MUST be identical in every city: deciding
        # it per city let a column be ranked in one and z-scored in the other,
        # putting them on different scales and making the two cities perfectly
        # separable (domain AUC 1.000 on `n_observed_years`). So it is computed
        # once on the design city and passed in here for the target.
        self.rankable = rankable

    def fit(self, X):
        X = np.asarray(X, dtype=np.float64)
        if self.mode is None:
            return self
        self.med_ = np.median(X, axis=0)
        q75, q25 = np.percentile(X, [75, 25], axis=0)
        self.iqr_ = np.where((q75 - q25) > 1e-12, q75 - q25, 1.0)
        if self.mode == "rank":
            self.q_ = np.linspace(0, 100, self.n_grid)
            self.grid_ = np.percentile(X, self.q_, axis=0)      # (n_grid, d)
            # Heavily tied columns must NOT be rank-transformed.
            #
            # Measured on the real data: `first_observed_year` is 1984 for
            # almost every pixel, and its domain AUC went from 0.507 raw to
            # 0.997 after ranking - the transform CREATED a perfect city marker
            # where there was none. A dominant tie makes the ECDF one long
            # plateau, and the plateau sits at a slightly different quantile in
            # each city, so every pixel lands on a different side of it.
            #
            # The test that matters is the size of the biggest tie, not the
            # number of distinct values: that column has 20 distinct values and
            # still breaks. Offending columns fall back to robust z.
            if self.rankable is not None:
                self.rankable_ = np.asarray(self.rankable, dtype=bool)
            else:
                self.rankable_ = rankable_mask(
                    X, min_unique=self.min_unique,
                    max_tie_share=self.max_tie_share)
        return self

    def transform(self, X):
        X = np.asarray(X, dtype=np.float64)
        if self.mode is None:
            return X
        if self.mode != "rank":
            return (X - self.med_) / self.iqr_
        out = (X - self.med_) / self.iqr_
        for j in np.flatnonzero(self.rankable_):
            out[:, j] = np.interp(X[:, j], self.grid_[:, j], self.q_ / 100.0,
                                  left=0.0, right=1.0)
        return out

    def fit_transform(self, X):
        return self.fit(X).transform(X)


def coral(X_src, X_tgt, eps: float = 1e-6):
    """
    CORAL: whiten the source with its own covariance, then re-colour it with
    the target covariance, so the two feature clouds have the same second-order
    statistics. Uses target FEATURES only - no labels.
    Returns the transformed source.
    """
    Xs = np.asarray(X_src, dtype=np.float64)
    Xt = np.asarray(X_tgt, dtype=np.float64)
    ms, mt = Xs.mean(0), Xt.mean(0)
    Cs = np.cov(Xs - ms, rowvar=False) + eps * np.eye(Xs.shape[1])
    Ct = np.cov(Xt - mt, rowvar=False) + eps * np.eye(Xt.shape[1])

    def _half(C, power):
        w, V = np.linalg.eigh(C)
        w = np.clip(w, eps, None)
        return V @ np.diag(w ** power) @ V.T

    A = _half(Cs, -0.5) @ _half(Ct, 0.5)
    return (Xs - ms) @ A + mt


# --------------------------------------------------------------------------- #
# Adversarial validation: which features simply say "this is Madrid"?
# --------------------------------------------------------------------------- #
def adversarial_validation(X_src, X_tgt, feature_names, n_estimators: int = 300,
                           seed: int = SEED, max_rows: int = 30_000):
    """
    Train a classifier to tell the two cities apart. An AUC near 1.0 means the
    domains are trivially separable; the top-importance features are the ones
    carrying the city signature rather than the age signal.

    LightGBM rather than a random forest, and capped at `max_rows` per city:
    this is called several times over ~380 features, and the AUC it reports
    saturates long before 30k rows per city, so a forest here would dominate
    the notebook's runtime for no extra information.

    Returns (auc, DataFrame of importances sorted descending).
    """
    import lightgbm as lgb
    import pandas as pd
    from sklearn.model_selection import cross_val_predict
    from sklearn.metrics import roc_auc_score

    rng = np.random.default_rng(seed)

    def cap(A):
        A = np.asarray(A, dtype=np.float32)
        if len(A) > max_rows:
            A = A[rng.choice(len(A), max_rows, replace=False)]
        return A

    Xs, Xt = cap(X_src), cap(X_tgt)
    X = np.vstack([Xs, Xt])
    y = np.r_[np.zeros(len(Xs)), np.ones(len(Xt))]
    clf = lgb.LGBMClassifier(objective="binary", n_estimators=n_estimators,
                             learning_rate=0.05, num_leaves=31,
                             min_child_samples=50, colsample_bytree=0.7,
                             subsample=0.8, subsample_freq=1, max_bin=127,
                             n_jobs=-1, random_state=seed, verbosity=-1)
    p = cross_val_predict(clf, X, y, cv=3, method="predict_proba", n_jobs=1)[:, 1]
    auc = roc_auc_score(y, p)
    clf.fit(X, y)
    gain = clf.feature_importances_.astype(float)
    imp = (pd.DataFrame({"feature": list(feature_names),
                         "city_importance": gain / max(gain.sum(), 1.0)})
           .sort_values("city_importance", ascending=False)
           .reset_index(drop=True))
    return float(auc), imp


# --------------------------------------------------------------------------- #
# Macro-F1 eval metrics
# --------------------------------------------------------------------------- #
def lgb_macro_f1(y_true, y_pred):
    """lightgbm 4.x feval. y_pred arrives as (n*4,) or (n,4) depending on version."""
    p = np.asarray(y_pred)
    if p.ndim == 1:
        p = p.reshape(len(y_true), -1, order="F") if p.size != len(y_true) else p
    lab = np.argmax(p, axis=1) if p.ndim == 2 else p
    return "macro_f1", float(f1_score(y_true, lab, average="macro", zero_division=0)), True


class CatMacroF1:
    """CatBoost custom eval metric (higher is better)."""

    def is_max_optimal(self):
        return True

    def evaluate(self, approxes, target, weight):
        p = np.array(approxes).T
        lab = np.argmax(p, axis=1)
        return float(f1_score(np.asarray(target, dtype=int), lab,
                              average="macro", zero_division=0)), 1.0

    def get_final_error(self, error, weight):
        return error


# --------------------------------------------------------------------------- #
# Default parameter sets.  Tuned toward robustness, not toward the last
# 0.002 of Madrid CV: the test set is a different city.
# --------------------------------------------------------------------------- #
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,
)

CATBOOST_PARAMS = dict(
    loss_function="MultiClass", classes_count=4, iterations=2000,
    learning_rate=0.03, depth=6, l2_leaf_reg=6.0, random_strength=1.5,
    bootstrap_type="Bernoulli", subsample=0.8, rsm=0.7,
    boosting_type="Plain", random_seed=SEED, verbose=False,
    allow_writing_files=False,
)


def fit_lgbm(X, y, params=None, class_weight=None, sample_weight=None):
    import lightgbm as lgb
    p = dict(LGBM_PARAMS)
    if params:
        p.update(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 fit_catboost(X, y, params=None, class_weights=None, sample_weight=None):
    from catboost import CatBoostClassifier
    p = dict(CATBOOST_PARAMS)
    if params:
        p.update(params)
    if class_weights is not None:
        p["class_weights"] = list(class_weights)
    m = CatBoostClassifier(**p)
    m.fit(np.asarray(X), np.asarray(y) - 1, sample_weight=sample_weight)
    return m


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


# --------------------------------------------------------------------------- #
# Ordinal decomposition: the four classes are ordered age bins, so three
# cumulative binary problems ("older than X?") use that structure.
# --------------------------------------------------------------------------- #
class OrdinalGBDT:
    """P(y > k) for k = 1,2,3, then differenced into class probabilities."""

    def __init__(self, fit_fn=fit_lgbm, params=None):
        self.fit_fn = fit_fn
        self.params = params

    def fit(self, X, y, sample_weight=None):
        import lightgbm as lgb
        X, y = np.asarray(X), np.asarray(y)
        params = dict(LGBM_PARAMS)
        params.update(self.params or {})
        params.pop("num_class", None)          # binary problems, not 4-way
        params["objective"] = "binary"
        self.models_ = []
        for k in (1, 2, 3):
            m = lgb.LGBMClassifier(**params)
            m.fit(X, (y > k).astype(int), sample_weight=sample_weight)
            self.models_.append(m)
        return self

    def predict_proba(self, X):
        X = np.asarray(X)
        s = np.column_stack([m.predict_proba(X)[:, 1] for m in self.models_])
        s = np.minimum.accumulate(s, axis=1)            # enforce monotonicity
        p1 = 1.0 - s[:, 0]
        p2 = s[:, 0] - s[:, 1]
        p3 = s[:, 1] - s[:, 2]
        p4 = s[:, 2]
        P = np.clip(np.column_stack([p1, p2, p3, p4]), 1e-9, None)
        return P / P.sum(axis=1, keepdims=True)


# --------------------------------------------------------------------------- #
# Calibration on the few-shot sample
# --------------------------------------------------------------------------- #
def stratification_weights(y_shot, target_prior):
    """
    The organisers hand us n labelled pixels *per class*, so the labelled sample
    is balanced while the city is not. Anything fitted on it therefore inherits
    a uniform class prior - which, in testing, cost class 4 almost 0.17 F1.

    These weights undo the stratification: each row is weighted by
    target_prior[c] / (1/4), so the weighted sample behaves like a random
    sample of the city. `target_prior` must come from somewhere legal - the EM
    estimate on unlabelled Amsterdam probabilities, or the organisers' own
    published class shares.
    """
    y = np.asarray(y_shot, dtype=int)
    tp = np.asarray(target_prior, dtype=float)
    tp = tp / tp.sum()
    counts = np.array([(y == c).sum() for c in CLASSES], dtype=float)
    share = counts / counts.sum()
    w = (tp / np.clip(share, 1e-9, None))[y - 1]
    return w * len(y) / w.sum()


def vector_scaling(proba_shot, y_shot, l2: float = 1.0, n_iter: int = 400,
                   lr: float = 0.1, sample_weight=None):
    """
    Learn per-class scale a_k and bias b_k on the log-probabilities using only
    the n-per-class labelled target sample:  q = softmax(a * log p + b).
    Eight parameters, so it survives n=5 per class.

    Pass `sample_weight=stratification_weights(...)` unless you really want the
    balanced prior of the shot sample baked into the output.
    Returns a callable proba -> recalibrated proba.
    """
    P = np.log(np.clip(np.asarray(proba_shot, dtype=np.float64), 1e-9, 1))
    y = np.asarray(y_shot, dtype=int) - 1
    Y = np.eye(4)[y]
    sw = (np.ones(len(y)) if sample_weight is None
          else np.asarray(sample_weight, dtype=float))
    sw = sw / sw.sum()
    a, b = np.ones(4), np.zeros(4)
    for _ in range(n_iter):
        z = P * a + b
        z -= z.max(axis=1, keepdims=True)
        q = np.exp(z)
        q /= q.sum(axis=1, keepdims=True)
        d = (q - Y) * sw[:, None]
        ga = (d * P).sum(0) + l2 * (a - 1.0)
        gb = d.sum(0) + l2 * b
        a -= lr * ga
        b -= lr * gb

    def apply(proba):
        lp = np.log(np.clip(np.asarray(proba, dtype=np.float64), 1e-9, 1))
        z = lp * a + b
        z -= z.max(axis=1, keepdims=True)
        e = np.exp(z)
        return e / e.sum(axis=1, keepdims=True)

    apply.params = (a, b)
    return apply


# --------------------------------------------------------------------------- #
# modelCreate adapters for evaluation.evaluate_protocol
# --------------------------------------------------------------------------- #
#: What `modelCreate` is allowed to see of the target city.
#
#: Prof. Thron's specification says Amsterdam data "may only be used to
#: EVALUATE modelCreate", and that every design decision must be made without
#: looking at Amsterdam. That rules out picking features or hyperparameters by
#: Amsterdam's score - which we never do. It leaves one genuine question open:
#: may modelCreate, as a FUNCTION applied to a new city, read that city's
#: unlabelled feature matrix at adaptation time?
#:
#: Both readings are defensible, so both are supported and both get reported.
#:
#:   "transductive" - modelCreate(D^(N)_city, X_city). Per-city statistics come
#:       from all the city's unlabelled pixels. This is standard unsupervised
#:       domain adaptation and matches the deployment story: pointing this at
#:       Baghdad, you have every Landsat pixel and merely lack labels. The
#:       organisers' own Notebook 3 is transductive in spirit, though it
#:       normalises Amsterdam with Madrid's scaler.
#:
#:   "inductive" - modelCreate(D^(N)_city) only. Per-city statistics are
#:       estimated from the N-per-class labelled sample alone. Strictly
#:       compliant on the tightest reading. At N=200 that is 800 pixels, which
#:       estimates a percentile well; at N=5 it is 20 and noticeably noisier.
#:
#: Ask the organisers which they intend. Until then, lead with whichever is
#: weaker and quote the other alongside it.
TARGET_ACCESS = ("transductive", "inductive")


def _target_reference(X_tgt_all, X_shot, access: str):
    """The rows a per-city statistic may be estimated from."""
    if access == "inductive":
        return np.asarray(X_shot, dtype=float)
    return np.asarray(X_tgt_all, dtype=float)


#: Default normalisation for the TREE models: none.
#:
#: Decided on Madrid, with blocked CV in which each held-out block is
#: normalised with its own statistics exactly as a new city would be:
#:   none 0.5989 +- 0.0230 | rank 0.5697 +- 0.0789 | robust z 0.5690 +- 0.0830
#: and the same ordering then held on Amsterdam (0.5928 against 0.5268).
#:
#: The reason is that the absolute level is signal, not only nuisance: building
#: era is spatially clustered, so a district's average brightness says something
#: real about when it was built, and normalising a region to its own median
#: subtracts that along with the city offset.
#:
#: DISTANCE-based methods are different and keep their normalisation. A nearest
#: neighbour over unscaled features is dominated by whichever column has the
#: largest units, which is an artefact of units rather than a statement about
#: transfer.
TREE_NORMALISE = None
DISTANCE_NORMALISE = "robustz"


def mc_zero_shot(fit=fit_lgbm, normalise: str | None = TREE_NORMALISE,
                 use_coral: bool = False, balanced: bool = True,
                 access: str = "transductive"):
    """
    Train on the source city. Under "transductive" access the few-shot sample
    is ignored entirely; under "inductive" it is still used, but only to
    estimate the target city's feature distribution, never its labels.
    """
    cache = {}

    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        ref = _target_reference(X_tgt_all, X_shot, access)
        rk = rankable_mask(X_src) if normalise == "rank" else None
        nt = CityNormaliser(normalise, rankable=rk).fit(ref) if normalise else None
        key = (normalise, use_coral, balanced, access)
        if key not in cache:
            ns = (CityNormaliser(normalise, rankable=rk).fit(X_src)
                  if normalise else None)
            Xs = ns.transform(X_src) if ns else np.asarray(X_src, float)
            if use_coral:
                Xs = coral(Xs, nt.transform(ref) if nt else ref)
            cw = balanced_class_weights(y_src) if balanced else None
            model = (fit(Xs, y_src, class_weight=dict(zip(range(4), cw)) if cw is not None else None)
                     if fit is fit_lgbm else fit(Xs, y_src, class_weights=cw))
            cache[key] = model
        model = cache[key]

        def predict(X):
            return model.predict_proba(nt.transform(X) if nt else X)
        return predict
    return model_create


def mc_recalibrated(fit=fit_lgbm, normalise: str | None = TREE_NORMALISE,
                    use_coral: bool = False, balanced: bool = True,
                    l2: float = 1.0, correct_stratification: bool = True,
                    tune_decision: bool = True):
    """
    Zero-shot model + vector scaling learned on the n-per-class sample, then a
    macro-F1 decision rule tuned on the same sample.

    `correct_stratification` reweights the shot sample towards the class prior
    that EM estimates from the unlabelled target features, so neither the
    calibration nor the decision rule inherits the artificial 25/25/25/25 prior
    of the sample.

    `tune_decision` matters more than the calibration: a well-calibrated model
    maximises log-likelihood, and argmax of a calibrated posterior is NOT the
    macro-F1 optimum. The shrinkage is tied to the sample size, so n=5 barely
    moves the rule and n=200 moves it freely.
    """
    base = mc_zero_shot(fit, normalise, use_coral, balanced)

    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        from evaluation import em_target_prior, tune_class_weights
        predict0 = base(X_src, y_src, X_tgt_all, X_shot, y_shot)
        sw = None
        if correct_stratification:
            src_prior = np.array([(np.asarray(y_src) == c).mean() for c in CLASSES])
            tgt_prior = em_target_prior(predict0(X_tgt_all), src_prior)
            sw = stratification_weights(y_shot, tgt_prior)
        p_shot = predict0(X_shot)
        cal = vector_scaling(p_shot, y_shot, l2=l2, sample_weight=sw)
        w = np.ones(4)
        if tune_decision:
            n_shot = len(y_shot)
            shrink = float(np.clip(1.0 - n_shot / 400.0, 0.0, 0.9))
            w = tune_class_weights(cal(p_shot), y_shot, sample_weight=sw,
                                   shrink=shrink)

        def predict(X):
            p = cal(predict0(X)) * w[None, :]
            return p / np.clip(p.sum(axis=1, keepdims=True), 1e-12, None)
        return predict
    return model_create


def mc_shot_weighted(fit=fit_lgbm, normalise: str | None = TREE_NORMALISE,
                     shot_weight: float = 30.0, balanced: bool = True):
    """
    Retrain on source + the few target rows, with the target rows upweighted.
    `shot_weight` must be chosen on the source city (e.g. by simulating a
    held-out Madrid district as a pseudo-target), never on the target score.
    """
    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        rk = rankable_mask(X_src) if normalise == "rank" else None
        ns = CityNormaliser(normalise, rankable=rk).fit(X_src) if normalise else None
        nt = (CityNormaliser(normalise, rankable=rk).fit(X_tgt_all)
              if normalise else None)
        Xs = ns.transform(X_src) if ns else np.asarray(X_src)
        Xsh = nt.transform(X_shot) if nt else np.asarray(X_shot)
        X = np.vstack([Xs, Xsh])
        y = np.r_[np.asarray(y_src), np.asarray(y_shot)]
        w = np.r_[np.ones(len(y_src)), np.full(len(y_shot), shot_weight)]
        if balanced:
            cw = balanced_class_weights(y)
            w = w * cw[np.asarray(y) - 1]
        model = fit(X, y, sample_weight=w)

        def predict(X_):
            return model.predict_proba(nt.transform(X_) if nt else X_)
        return predict
    return model_create


def mc_knn(normalise: str = DISTANCE_NORMALISE, k: int | None = None,
           columns=None):
    """
    The pure-retrieval classifier: distance-weighted k-NN over the n labelled
    target pixels alone, ignoring the source city completely.

    Measure this FIRST. It costs seconds and it bounds everything else: if
    50-per-class k-NN already beats the reported 0.34 zero-shot Random Forest,
    then the retrieval layer is the story and the source model is support, not
    the other way round. It is also the baseline the retrieval features in the
    LLM prompt have to beat to justify their complexity.
    """
    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        rk = rankable_mask(X_src) if normalise == "rank" else None
        nt = CityNormaliser(normalise, rankable=rk).fit(X_tgt_all)
        Xs = nt.transform(X_shot)
        y_shot = np.asarray(y_shot)
        kk = k if k is not None else int(min(15, max(3, 3 * len(y_shot) // 4)))

        # Squared distances via the (a-b)^2 = a^2 - 2ab + b^2 identity, in row
        # chunks. The naive Xn[:, None, :] - Xs[None, :, :] would allocate
        # n_query x n_pool x n_features floats - 63 GB for 26k queries against
        # 800 labelled pixels over 380 features.
        pool_sq = (Xs ** 2).sum(1)
        CHUNK = 2048

        def predict(X):
            Xn = nt.transform(X)
            P = np.zeros((len(Xn), 4))
            for s in range(0, len(Xn), CHUNK):
                Q = Xn[s:s + CHUNK]
                d2 = (Q ** 2).sum(1)[:, None] - 2.0 * (Q @ Xs.T) + pool_sq[None, :]
                d = np.sqrt(np.maximum(d2, 0.0))
                idx = np.argpartition(d, kk - 1, axis=1)[:, :kk]
                dd = np.take_along_axis(d, idx, axis=1)
                tau = np.median(dd, axis=1, keepdims=True) + 1e-9
                w = np.exp(-(dd / tau) ** 2)
                lab = y_shot[idx]
                for j, c in enumerate(CLASSES):
                    P[s:s + CHUNK, j] = (w * (lab == c)).sum(axis=1)
            P = np.clip(P, 1e-9, None)
            return P / P.sum(axis=1, keepdims=True)
        return predict
    return model_create


def mc_refit(normalise: str | None = TREE_NORMALISE, decay: float = 0.9,
             balanced: bool = False):
    """
    LightGBM `Booster.refit`: keep the tree STRUCTURE learned on the source city
    and move only the leaf values towards the target, blended by `decay`.

    This is the right shape for the problem. Tree structure is what needs tens
    of thousands of rows; leaf values are a handful of numbers per tree and can
    be nudged by a few hundred. Recommended from n=25 upwards; below that there
    is nothing estimable and plain recalibration is safer.
    """
    cache = {}

    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        key = (normalise, balanced)
        if key not in cache:
            rk = rankable_mask(X_src) if normalise == "rank" else None
            ns = CityNormaliser(normalise, rankable=rk).fit(X_src)
            nt = CityNormaliser(normalise, rankable=rk).fit(X_tgt_all)
            cw = (dict(zip(range(4), balanced_class_weights(y_src)))
                  if balanced else None)
            cache[key] = (fit_lgbm(ns.transform(X_src), y_src, class_weight=cw), nt)
        model, nt = cache[key]
        Xsh = nt.transform(X_shot)
        booster = model.booster_.refit(Xsh, np.asarray(y_shot) - 1,
                                       decay_rate=decay)

        def predict(X):
            P = booster.predict(nt.transform(X))
            P = np.asarray(P, dtype=float)
            if P.ndim == 1:
                P = P.reshape(-1, 4)
            return P / np.clip(P.sum(axis=1, keepdims=True), 1e-12, None)
        return predict
    return model_create


def single_feature_domain_auc(X_src, X_tgt, feature_names):
    """
    One AUC per feature for the question "which city is this?".

    0.5 means the feature is domain-invariant; 1.0 means it is a pure city
    marker and a Madrid-trained split on it is meaningless in Amsterdam. Run it
    before and after per-city normalisation: the table of the two columns is the
    most convincing single figure for the write-up, and it is the honest basis
    for dropping features. Prune conservatively - drop only a feature that is
    still above ~0.90 after normalisation AND carries little Madrid importance.
    """
    import pandas as pd
    from sklearn.metrics import roc_auc_score

    Xs = np.asarray(X_src, dtype=float)
    Xt = np.asarray(X_tgt, dtype=float)
    y = np.r_[np.zeros(len(Xs)), np.ones(len(Xt))]
    rows = []
    for j, name in enumerate(feature_names):
        v = np.r_[Xs[:, j], Xt[:, j]]
        ok = np.isfinite(v)
        if ok.sum() < 10 or np.std(v[ok]) == 0:
            rows.append(dict(feature=name, domain_auc=0.5))
            continue
        a = roc_auc_score(y[ok], v[ok])
        rows.append(dict(feature=name, domain_auc=max(a, 1 - a)))
    return (pd.DataFrame(rows).sort_values("domain_auc", ascending=False)
            .reset_index(drop=True))


def mc_prototype_official():
    """
    The organisers' own few-shot baseline, reproduced exactly: class means of
    the labelled target sample, nearest prototype by plain Euclidean distance,
    on features standardised with the SOURCE city's scaler.

    This is the row every other row in the table has to beat, and it is a
    stronger baseline than it looks. Both the prototypes and the queries come
    from the target city and receive the same transform, so any affine
    difference between the cities cancels exactly in a nearest-prototype
    comparison. The source scaler therefore does no harm here; it only sets the
    relative weighting of features. In synthetic testing this plain version beat
    a rank-normalised Mahalanobis version by 0.05 macro F1 at n=25, because the
    rank transform distorts the cluster geometry that prototypes depend on.
    Beat it on measurements, not on sophistication.
    """
    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        mu_s = np.mean(X_src, axis=0)
        sd_s = np.std(X_src, axis=0)
        sd_s = np.where(sd_s > 1e-12, sd_s, 1.0)
        Xsh = (np.asarray(X_shot, float) - mu_s) / sd_s
        y_shot = np.asarray(y_shot)
        present = [c for c in CLASSES if (y_shot == c).sum() > 0]
        proto = np.stack([Xsh[y_shot == c].mean(0) for c in present])

        def predict(X):
            Xn = (np.asarray(X, float) - mu_s) / sd_s
            d = np.linalg.norm(Xn[:, None, :] - proto[None, :, :], axis=2)
            P = np.zeros((len(Xn), 4))
            # hard nearest-prototype assignment, as in their code
            for row, j in enumerate(np.argmin(d, axis=1)):
                P[row, present[j] - 1] = 1.0
            return np.clip(P, 1e-9, None) / np.clip(P, 1e-9, None).sum(1, keepdims=True)
        return predict
    return model_create


def mc_prototype(normalise: str | None = DISTANCE_NORMALISE,
                 metric: str = "euclidean",
                 shrink: float = 0.5, temperature: float = 1.0):
    """
    Prototype classifier: one class mean from the n labelled target pixels,
    nearest prototype wins. Needs no source model, so it is the honest floor.

    `normalise` defaults to per-city robust z rather than rank. Rank mapping
    forces each marginal to be uniform, which is excellent for tree splits and
    bad here - it distorts the cluster shape that a distance to a class mean
    depends on. Measured: rank + Mahalanobis lost 0.05 macro F1 at n=25 against
    the plain version on synthetic data.

    `metric`:
      "euclidean"  - isotropic distance.
      "within"     - Mahalanobis on the POOLED WITHIN-CLASS covariance of the
                     shots (shrunk towards diagonal). This is the statistically
                     right metric for prototypes, and it is not the same thing
                     as using the target city's total covariance, which is
                     dominated by the BETWEEN-class spread and therefore shrinks
                     exactly the directions that separate the classes.
      "total"      - Mahalanobis on the target city's total covariance; kept
                     only so the ablation table can show why it is worse.
    """
    def model_create(X_src, y_src, X_tgt_all, X_shot, y_shot):
        rk = rankable_mask(X_src) if normalise == "rank" else None
        nt = (CityNormaliser(normalise, rankable=rk).fit(X_tgt_all)
              if normalise else None)
        Xt = nt.transform(X_tgt_all) if nt else np.asarray(X_tgt_all, float)
        Xsh = nt.transform(X_shot) if nt else np.asarray(X_shot, float)
        y_shot = np.asarray(y_shot)
        present = [c for c in CLASSES if (y_shot == c).sum() > 0]
        mu = np.stack([Xsh[y_shot == c].mean(0) for c in present])
        d_ = Xt.shape[1]

        if metric == "within":
            # Pooled within-class scatter of the shots. Macro F1 on synthetic
            # anisotropic data against the organisers' Euclidean prototype:
            # n=5 0.719 vs 0.649, n=25 0.967 vs 0.821, n=100 0.988 vs 0.833 -
            # the largest single gain of any adaptation step tried.
            #
            # Without the blending below it LOSES at n=5, where 20 rows cannot
            # fill a 20x20 covariance, so the estimate is pulled towards the
            # target city's total covariance and towards its own diagonal by how
            # much data there actually is.
            #
            # Honest caveat: those synthetic classes are Gaussian sharing one
            # anisotropic within-class covariance, i.e. exactly the model this
            # metric assumes. Expect less on real Landsat features; trust the
            # ORDERING of the four variants, not the absolute numbers.
            parts = [Xsh[y_shot == c] - Xsh[y_shot == c].mean(0)
                     for c in present if (y_shot == c).sum() > 1]
            resid = np.vstack(parts) if parts else np.zeros((0, d_))
            n_res = len(resid)
            Cw = (np.cov(resid, rowvar=False) if n_res > d_
                  else (np.diag(np.var(resid, axis=0) + 1e-6) if n_res
                        else np.eye(d_)))
            Ct = np.cov(Xt, rowvar=False)
            a = float(np.clip(d_ / (d_ + max(n_res, 1)), 0.0, 0.8))
            C = (1 - a) * Cw + a * Ct
            s = float(np.clip(shrink * d_ / (d_ + max(n_res, 1)) * 2.0, 0.1, 0.9))
            C = (1 - s) * C + s * np.diag(np.diag(C) + 1e-6) + 1e-6 * np.eye(d_)
            A = np.linalg.pinv(C)
        elif metric == "total":
            Ct = np.cov(Xt, rowvar=False)
            Ct = (1 - shrink) * Ct + shrink * np.diag(np.diag(Ct)) + 1e-6 * np.eye(d_)
            A = np.linalg.pinv(Ct)
        else:
            A = np.eye(d_)

        def predict(X):
            Xn = nt.transform(X) if nt else np.asarray(X, float)
            d = np.stack([np.einsum("ij,jk,ik->i", Xn - m, A, Xn - m) for m in mu],
                         axis=1)
            z = -0.5 * d / max(temperature, 1e-6)
            z -= z.max(axis=1, keepdims=True)
            e = np.exp(z)
            e = e / e.sum(axis=1, keepdims=True)
            P = np.full((len(Xn), 4), 1e-9)
            for j, c in enumerate(present):
                P[:, c - 1] = e[:, j]
            return P / P.sum(axis=1, keepdims=True)
        return predict
    return model_create