"""
Raw wide-format Landsat parquet  ->  one feature row per 30 m pixel.

Two feature sets come out of here:

  * the OFFICIAL 60  - a faithful re-implementation of the organisers'
    Notebook 3, so our numbers stay comparable with the published baseline;
  * the EXTENDED set - physically scaled surface reflectance, correct spectral
    indices, construction-event ("step") detection, per-decade means, trends
    and early/late contrasts. This is what the LLM verbaliser and the strong
    GBDT models actually use.

Everything is vectorised over a (n_pixels x n_years) matrix per band, so a whole
city is processed in a couple of minutes on a laptop CPU.
"""
from __future__ import annotations

import warnings

import numpy as np
import pandas as pd

from config import (BASE_BANDS, BLUE_MAX, CLASS_EVENTS, EARLY_CUTOFF,
                    FEATURE_COLS_60, INDEX_COLS, LATE_START, PID, SR_OFFSET,
                    SR_SCALE)

EPS = 1e-6


# --------------------------------------------------------------------------- #
# Stage 1 - load, label, flatten, gap-fill  (official logic)
# --------------------------------------------------------------------------- #
def assign_age_class(df: pd.DataFrame) -> pd.DataFrame:
    """Bin weighted_mean_year into classes 1-4 with city-specific boundaries."""
    out = df
    out["age_class"] = pd.Series(pd.NA, index=out.index, dtype="Int64")
    for city, ev in CLASS_EVENTS.items():
        mask = (out["city"] == city).to_numpy()
        if mask.any():
            binned = pd.cut(out.loc[mask, "weighted_mean_year"],
                            bins=[-np.inf, ev, 1984, 2004, np.inf],
                            labels=[1, 2, 3, 4])
            out.loc[mask, "age_class"] = binned.astype("Int64")
    return out


def make_flat_view(df: pd.DataFrame, agg: str = "first") -> pd.DataFrame:
    """
    One value per (pixel, year) per band; drop haze-contaminated rows.

    agg="first"  : the first QA-valid observation, exactly as the organisers do.
    agg="median" : the median of all QA-valid observations in that year. With up
                   to three summer scenes per year, a single undetected thin
                   cloud moves a yearly value by more than any real surface
                   change, so the median is the more robust choice - use it for
                   our own features and keep "first" to reproduce the official
                   60 exactly.
    """
    keep = [c for c in ["city", "year", "px_key", "py_key", "coverage",
                        "weighted_mean_year", "age_class"] if c in df.columns]
    out = df[keep].copy()
    n = len(df)

    # Which satellite supplied the observation we end up using.
    #
    # This matters more than it looks. Reading scene_id (the dataset_* string
    # column disagrees with it and is not to be trusted), the satellite chosen
    # for a given year DIFFERS BETWEEN THE TWO CITIES in 11 of the 25 years
    # from 2001 on: Madrid's late record is largely Landsat 7 ETM+, which has
    # been SLC-off and striped since 2003, while Amsterdam's is largely
    # Landsat 8/9 OLI. So part of any early-vs-late or year-to-year feature
    # measures which instrument flew over, not what the building did - and that
    # part is city-specific, which is the worst kind of nuisance here. Carrying
    # the sensor through lets a model condition on it instead of confounding
    # through it, and lets us measure how much it costs.
    sat_code = np.zeros(n, dtype=np.int8)        # 0 unknown, 1 TM, 2 ETM+, 3 OLI
    _MAP = {"LT04": 1, "LT05": 1, "LE07": 2, "LC08": 3, "LC09": 3}
    assigned = np.zeros(n, dtype=bool)
    for slot in (1, 2, 3):
        sc, qc = f"scene_id_{slot}", f"qa_valid_{slot}"
        if sc not in df.columns or qc not in df.columns:
            continue
        valid = df[qc].fillna(False).to_numpy(dtype=bool)
        codes = (df[sc].astype("string").str.slice(0, 4)
                 .map(_MAP).fillna(0).to_numpy(dtype=np.int8))
        take = valid & ~assigned & (codes > 0)
        sat_code[take] = codes[take]
        assigned |= take
    out["sat_code"] = sat_code

    for band in BASE_BANDS:
        if agg == "median":
            stack = []
            for slot in (1, 2, 3):
                bc, qc = f"{band}_{slot}", f"qa_valid_{slot}"
                if bc not in df.columns or qc not in df.columns:
                    continue
                v = np.asarray(df[bc], dtype=np.float32).copy()
                ok = df[qc].fillna(False).to_numpy(dtype=bool)
                v[~ok] = np.nan
                stack.append(v)
            if stack:
                with warnings.catch_warnings():       # all-NaN rows are expected
                    warnings.simplefilter("ignore", RuntimeWarning)
                    col = np.nanmedian(np.stack(stack), axis=0).astype(np.float32)
            else:
                col = np.full(n, np.nan, dtype=np.float32)
        else:
            col = np.full(n, np.nan, dtype=np.float32)
            for slot in (1, 2, 3):
                bc, qc = f"{band}_{slot}", f"qa_valid_{slot}"
                if bc not in df.columns or qc not in df.columns:
                    continue
                valid = df[qc].fillna(False).to_numpy(dtype=bool)
                vals = np.asarray(df[bc], dtype=np.float32)
                take = np.isnan(col) & valid & ~np.isnan(vals)
                col[take] = vals[take]
        out[band] = col
    out = out.dropna(subset=BASE_BANDS)
    before = len(out)
    out = out[out["Blue"] <= BLUE_MAX]
    removed = before - len(out)
    if removed:
        print(f"  removed {removed:,} haze rows (Blue > {BLUE_MAX:,})")
    return out


def period_coverage(flat: pd.DataFrame) -> pd.DataFrame:
    """has_early_data / has_late_data, computed on OBSERVED years (pre gap-fill)."""
    pixels = flat[PID].drop_duplicates()
    early = flat.loc[flat["year"] <= EARLY_CUTOFF, PID].drop_duplicates().assign(has_early_data=1.0)
    late = flat.loc[flat["year"] >= LATE_START, PID].drop_duplicates().assign(has_late_data=1.0)
    return (pixels.merge(early, on=PID, how="left")
                  .merge(late, on=PID, how="left")
                  .fillna({"has_early_data": 0.0, "has_late_data": 0.0}))


def to_year_matrices(flat: pd.DataFrame):
    """
    Pivot the observed rows into one (n_pixels x n_years) matrix per band and
    gap-fill exactly as the organisers do: linear interior, nearest at the edges.

    Returns (pixel index DataFrame, years array, {band: matrix}, observed mask)
    """
    years = np.arange(int(flat["year"].min()), int(flat["year"].max()) + 1)
    pix = flat[PID].drop_duplicates().sort_values(PID).reset_index(drop=True)
    pix_pos = pd.Series(np.arange(len(pix)), index=pd.MultiIndex.from_frame(pix))

    ri = pix_pos.loc[pd.MultiIndex.from_frame(flat[PID])].to_numpy()
    ci = (flat["year"].to_numpy() - years[0]).astype(np.int32)

    mats = {}
    observed = np.zeros((len(pix), len(years)), dtype=bool)
    observed[ri, ci] = True
    sat = np.zeros((len(pix), len(years)), dtype=np.int8)
    if "sat_code" in flat.columns:
        sat[ri, ci] = flat["sat_code"].to_numpy(dtype=np.int8)
    for band in BASE_BANDS:
        M = np.full((len(pix), len(years)), np.nan, dtype=np.float32)
        M[ri, ci] = flat[band].to_numpy(dtype=np.float32)
        M = (pd.DataFrame(M).interpolate(axis=1, method="linear")
                            .bfill(axis=1).ffill(axis=1).to_numpy(dtype=np.float32))
        mats[band] = M
    return pix, years, mats, observed, sat


# --------------------------------------------------------------------------- #
# Stage 2 - the official 60 features, from DN values
# --------------------------------------------------------------------------- #
def _indices_from(mats: dict) -> dict:
    """The organisers' five indices, exactly as their Notebook 3 defines them."""
    B, G, R = mats["Blue"], mats["Green"], mats["Red"]
    N, S1, S2 = mats["NIR"], mats["SWIR1"], mats["SWIR2"]
    return {
        "NDVI":  (N - R) / (N + R + EPS),
        "NDBI":  (S1 - N) / (S1 + N + EPS),
        "UI":    (S2 - N) / (S2 + N + EPS),
        "MNDWI": (G - S1) / (G + S1 + EPS),
        "BSI":   ((S1 + R) - (N + B)) / ((S1 + R) + (N + B) + EPS),
    }


# Tasseled-cap coefficients for OLI surface reflectance (Baig et al. 2014).
# Wetness in particular separates impervious surfaces from soil better than
# any of the five indices above.
TC_COEF = {
    "TCB": (0.3029, 0.2786, 0.4733, 0.5599, 0.5080, 0.1872),
    "TCG": (-0.2941, -0.2430, -0.5424, 0.7276, 0.0713, -0.1608),
    "TCW": (0.1511, 0.1973, 0.3283, 0.3407, -0.7117, -0.4559),
}


def _extra_indices_from(sr: dict) -> dict:
    """NBR, tasseled cap, and chromaticity (band / total brightness).

    Chromaticity matters for transfer: dividing each band by the pixel's own
    6-band sum removes albedo, illumination and the city-level brightness
    offset entirely, and keeps only the spectral SHAPE.
    """
    order = [sr[b] for b in BASE_BANDS]
    total = sum(order) + EPS
    out = {"NBR": (sr["NIR"] - sr["SWIR2"]) / (sr["NIR"] + sr["SWIR2"] + EPS)}
    for name, coef in TC_COEF.items():
        out[name] = sum(c * m for c, m in zip(coef, order))
    for b in BASE_BANDS:
        out[f"chrom{b}"] = sr[b] / total
    return out


def _median_filter3(M: np.ndarray) -> np.ndarray:
    """
    3-point running median along the time axis.

    One undetected thin cloud moves a yearly value by more than any real
    surface change, so the change detector runs on the filtered series while
    the reported statistics stay on the raw one.
    """
    a = np.pad(M, ((0, 0), (1, 1)), mode="edge")
    return np.median(np.stack([a[:, :-2], a[:, 1:-1], a[:, 2:]]), axis=0)


def build_official_60(pix, years, mats, indicators) -> pd.DataFrame:
    """Reproduce the organisers' 60 features (DN units, ddof=1 std)."""
    idx = _indices_from(mats)
    early = years <= EARLY_CUTOFF
    late = years >= LATE_START
    out = {}
    for b in BASE_BANDS:
        M = mats[b]
        out[f"{b}_mean"] = M.mean(1)
        out[f"{b}_std"] = M.std(1, ddof=1)
        out[f"{b}_early_mean"] = M[:, early].mean(1)
        out[f"{b}_early_std"] = M[:, early].std(1, ddof=1)
        out[f"{b}_late_mean"] = M[:, late].mean(1)
        out[f"{b}_late_std"] = M[:, late].std(1, ddof=1)
        D = np.diff(M, axis=1)
        out[f"d_{b}_mean"] = D.mean(1)
        out[f"d_{b}_std"] = D.std(1, ddof=1)
    for i in INDEX_COLS:
        out[f"{i}_mean"] = idx[i].mean(1)
        out[f"{i}_std"] = idx[i].std(1, ddof=1)
    df = pd.concat([pix, pd.DataFrame(out, index=pix.index)], axis=1)
    df = df.merge(indicators, on=PID, how="left")
    return df[PID + FEATURE_COLS_60]


# --------------------------------------------------------------------------- #
# Stage 3 - extended, physically-scaled features
# --------------------------------------------------------------------------- #
DECADES = [(1984, 1993), (1994, 2003), (2004, 2013), (2014, 2024)]
DECADE_TAGS = ["84_93", "94_03", "04_13", "14_24"]
STEP_SIGNALS = ["brightness", "SWIR1", "NDVI"]


def _slope_per_decade(M: np.ndarray, years: np.ndarray) -> np.ndarray:
    """OLS slope of each row against year, expressed per decade."""
    t = (years - years.mean()).astype(np.float32)
    denom = float((t * t).sum())
    return (M * t).sum(1) / denom * 10.0


def _step_stats(M: np.ndarray, years: np.ndarray, win: int = 3):
    """
    Construction-event detector. For every candidate break year, compare the
    mean of the `win` years before with the mean of the `win` years after.
    Returns (year of the largest absolute jump, signed size of that jump,
    number of year-on-year jumps above 3 x MAD).
    """
    n, T = M.shape
    cs = np.concatenate([np.zeros((n, 1), np.float32),
                         np.cumsum(M, axis=1, dtype=np.float32)], axis=1)

    def win_mean(a, b):
        return (cs[:, b] - cs[:, a]) / np.float32(b - a)

    cuts = np.arange(win, T - win + 1)
    diffs = np.stack([win_mean(c, c + win) - win_mean(c - win, c) for c in cuts], axis=1)

    # A two-sample t statistic, not just the raw difference.
    #
    # argmax over candidate years ALWAYS returns something, so a pixel that was
    # never rebuilt still gets a "biggest step" - and inspection caught the
    # verbaliser telling the model that a building from 1924 changed abruptly in
    # 1998. What distinguishes a real construction event is the size of the step
    # RELATIVE TO the scatter within the two windows being compared, which is
    # exactly a t statistic. The earlier version divided by the scatter about
    # the whole-series mean, which the step itself inflates, so it could not
    # separate the two cases.
    cs2 = np.concatenate([np.zeros((n, 1), np.float32),
                          np.cumsum(M.astype(np.float32) ** 2, axis=1, dtype=np.float32)],
                         axis=1)

    def win_var(a, b):
        k = np.float32(b - a)
        mu = (cs[:, b] - cs[:, a]) / k
        ex2 = (cs2[:, b] - cs2[:, a]) / k
        return np.maximum(ex2 - mu * mu, 0.0) * (k / max(k - 1.0, 1.0))

    tstats = np.empty_like(diffs)
    for m_i, c in enumerate(cuts):
        vb, va = win_var(c - win, c), win_var(c, c + win)
        pooled = np.sqrt(np.maximum((vb + va) / 2.0, 0.0) * (2.0 / win)) + EPS
        tstats[:, m_i] = diffs[:, m_i] / pooled

    # The break year is chosen by the largest MAGNITUDE, not the largest t.
    #
    # Selecting on t was tried and was a disaster: Amsterdam's agreement with
    # the true construction year fell from 69.4% within five years to 31.5%,
    # and Madrid's from 44.6% to 28.3%. The t statistic is maximised wherever
    # the pooled variance of a 3-year window happens to be near zero, which is
    # noise, not a construction event. The magnitude is the signal.
    #
    # The t value at that same break is still reported as a feature, where it
    # means "how large is this shift relative to the local scatter" and the
    # trees can use it or ignore it. It is NOT used for any wording in the
    # prompt - see the note in verbalize.render_prompt.
    k = np.argmax(np.abs(diffs), axis=1)
    rows = np.arange(n)
    step_mag = diffs[rows, k]
    step_t = np.abs(tstats[rows, k])
    step_year = years[cuts[k]].astype(np.float32)

    D = np.diff(M, axis=1)
    mad = np.median(np.abs(D - np.median(D, axis=1, keepdims=True)), axis=1) + EPS
    n_jumps = (np.abs(D) > (3.0 * mad[:, None])).sum(1).astype(np.float32)
    return step_year, step_mag, n_jumps, step_t


SPIKE_SIGNALS = ["BSI", "SWIR1", "brightness"]
# Sensor eras. The early/late split at 2004 also straddles a change of
# atmospheric-correction algorithm (LEDAPS for TM/ETM+, LaSRC for OLI), and the
# artefact is latitude-dependent, so part of any early/late contrast measures
# which algorithm ran rather than what the building did. Splitting at the
# sensor boundary as well makes that confound explicit instead of hidden.
TM_ERA = (1984, 2011)
OLI_ERA = (2013, 2024)


SAT_NAMES = {1: "tm", 2: "etm", 3: "oli"}


def sensor_features(pix, years, observed, sat) -> dict:
    """
    Per-pixel sensor composition: which instrument actually supplied each year.

    Not a building property at all - which is the point. These columns let a
    model absorb the sensor nuisance explicitly instead of attributing it to
    building age, and they let us measure how much of the Madrid/Amsterdam gap
    is instrument rather than architecture.
    """
    early = years <= EARLY_CUTOFF
    late = years >= LATE_START
    out = {}
    n_obs = np.clip(observed.sum(1), 1, None).astype(np.float32)
    for code, nm in SAT_NAMES.items():
        hit = (sat == code)
        out[f"sat_share_{nm}"] = hit.sum(1) / n_obs
        out[f"sat_share_{nm}_early"] = (hit & early[None, :]).sum(1) / np.clip(
            (observed & early[None, :]).sum(1), 1, None)
        out[f"sat_share_{nm}_late"] = (hit & late[None, :]).sum(1) / np.clip(
            (observed & late[None, :]).sum(1), 1, None)
    # how often the instrument changes from one observed year to the next
    switches = np.zeros(len(sat), dtype=np.float32)
    for i in range(sat.shape[0]):
        s = sat[i][observed[i]]
        s = s[s > 0]
        if len(s) > 1:
            switches[i] = float((s[1:] != s[:-1]).sum())
    out["sat_switches"] = switches
    return {k: np.asarray(v, dtype=np.float32) for k, v in out.items()}


def build_extended(pix, years, mats_dn, observed, sat=None) -> pd.DataFrame:
    """Surface-reflectance features: correct indices, trends, steps, decades."""
    sr = {b: mats_dn[b] * np.float32(SR_SCALE) + np.float32(SR_OFFSET)
          for b in BASE_BANDS}
    idx = _indices_from(sr)
    extra = _extra_indices_from(sr)
    bright = np.mean(np.stack([sr[b] for b in BASE_BANDS]), axis=0, dtype=np.float32)
    signals = {"brightness": bright}
    signals.update(sr)
    signals.update(idx)
    signals.update(extra)

    early = years <= EARLY_CUTOFF
    late = years >= LATE_START
    tm = (years >= TM_ERA[0]) & (years <= TM_ERA[1])
    oli = (years >= OLI_ERA[0]) & (years <= OLI_ERA[1])
    out = {}

    # Full statistics for the signals worth the columns; a short set for the
    # chromaticity and tasseled-cap extras, to keep the table manageable.
    rich = {"brightness", *BASE_BANDS, *INDEX_COLS}
    for name, M in signals.items():
        out[f"sr_{name}_mean"] = M.mean(1)
        out[f"sr_{name}_std"] = M.std(1, ddof=1)
        out[f"sr_{name}_early_mean"] = M[:, early].mean(1)
        out[f"sr_{name}_late_mean"] = M[:, late].mean(1)
        out[f"sr_{name}_late_minus_early"] = (out[f"sr_{name}_late_mean"]
                                             - out[f"sr_{name}_early_mean"])
        # std scales with brightness, so a raw std is partly a city marker;
        # the coefficient of variation is not.
        out[f"sr_{name}_cv"] = out[f"sr_{name}_std"] / (np.abs(out[f"sr_{name}_mean"]) + EPS)
        if name not in rich:
            continue
        out[f"sr_{name}_early_std"] = M[:, early].std(1, ddof=1)
        out[f"sr_{name}_late_std"] = M[:, late].std(1, ddof=1)
        out[f"sr_{name}_slope_decade"] = _slope_per_decade(M, years)
        out[f"sr_{name}_min"] = M.min(1)
        out[f"sr_{name}_max"] = M.max(1)
        out[f"sr_{name}_tm_mean"] = M[:, tm].mean(1)
        out[f"sr_{name}_oli_mean"] = M[:, oli].mean(1)
        out[f"sr_{name}_oli_minus_tm"] = out[f"sr_{name}_oli_mean"] - out[f"sr_{name}_tm_mean"]
        for (y0, y1), tag in zip(DECADES, DECADE_TAGS):
            sel = (years >= y0) & (years <= y1)
            out[f"sr_{name}_dec{tag}_mean"] = M[:, sel].mean(1)

    # ---- construction events, on the cloud-filtered series -----------------
    smoothed = {k: _median_filter3(signals[k]) for k in set(STEP_SIGNALS) | set(SPIKE_SIGNALS)}
    for name in STEP_SIGNALS:
        sy, sm, nj, st = _step_stats(smoothed[name], years)
        out[f"step_{name}_year"] = sy
        out[f"step_{name}_mag"] = sm
        out[f"step_{name}_njumps"] = nj
        out[f"step_{name}_t"] = st
        resid = signals[name] - signals[name].mean(1, keepdims=True)
        out[f"step_{name}_snr"] = np.abs(sm) / (resid.std(1, ddof=1) + EPS)

    # ---- transient bare-soil spike: a building site is bare for 1-3 years --
    for name in SPIKE_SIGNALS:
        M = smoothed[name]
        med = np.median(M, axis=1)
        mad = np.median(np.abs(M - med[:, None]), axis=1) + EPS
        k = np.argmax(M, axis=1)
        out[f"spike_{name}_year"] = years[k].astype(np.float32)
        out[f"spike_{name}_z"] = (M[np.arange(len(M)), k] - med) / mad
        # does it come back down? a permanent rise is a new roof, a spike that
        # returns to baseline is a construction site
        after = np.clip(k + 3, 0, M.shape[1] - 1)
        out[f"spike_{name}_returns"] = ((M[np.arange(len(M)), after] - med) / mad)

    # ---- post-construction greening: street trees take decades to mature ---
    ndvi = signals["NDVI"]
    sy_b = out["step_brightness_year"]
    slope_after = np.zeros(len(ndvi), dtype=np.float32)
    years_since = np.zeros(len(ndvi), dtype=np.float32)
    for i in range(len(ndvi)):
        k = int(np.searchsorted(years, sy_b[i]))
        seg = ndvi[i, k:]
        years_since[i] = float(years[-1] - sy_b[i])
        if len(seg) >= 5:
            t = np.arange(len(seg), dtype=np.float32)
            t -= t.mean()
            slope_after[i] = float((seg * t).sum() / (t * t).sum() * 10.0)
    out["post_step_NDVI_slope_decade"] = slope_after
    out["years_since_step"] = years_since

    # ---- what was there before the change ----------------------------------
    # 0 vegetated, 1 bare, 2 already built, 3 water. For a pixel that changed
    # after 2004 this separates greenfield expansion from brownfield infill and
    # from land reclamation, which have completely different trajectories but
    # land in the same class.
    pre_ndvi = out["sr_NDVI_early_mean"]
    pre_mndwi = out["sr_MNDWI_early_mean"]
    pre_ndbi = out["sr_NDBI_early_mean"]
    pre = np.full(len(pre_ndvi), 2.0, dtype=np.float32)
    pre[pre_ndvi > 0.30] = 0.0
    pre[(pre_ndvi <= 0.30) & (pre_ndbi < 0.0) & (out["sr_BSI_early_mean"] > 0.05)] = 1.0
    pre[pre_mndwi > 0.10] = 3.0
    out["pre_state"] = pre

    if sat is not None:
        out.update(sensor_features(pix, years, observed, sat))

    out["n_observed_years"] = observed.sum(1).astype(np.float32)
    out["first_observed_year"] = years[np.argmax(observed, axis=1)].astype(np.float32)
    last = observed.shape[1] - 1 - np.argmax(observed[:, ::-1], axis=1)
    out["last_observed_year"] = years[last].astype(np.float32)

    return pd.concat([pix, pd.DataFrame(out, index=pix.index)], axis=1)


PRE_STATE_WORDS = {0: "vegetated - farmland, scrub or garden",
                   1: "bare soil or a cleared site",
                   2: "already built on or paved",
                   3: "open water (later reclaimed)"}


def add_city_geometry(df: pd.DataFrame) -> pd.DataFrame:
    """
    Two purely geometric, label-free features.

    `radial_frac`: distance from the city centroid, divided by the city's own
    95th percentile, so it is comparable between cities. European cities share a
    morphological gradient - oldest and densest in the historic core - and this
    is the only transferable handle on the class 1 / class 2 split, whose
    boundary year (1960 vs 1945) has no spectral correlate at all.

    `dist_to_water_px`: distance to the nearest pixel that looks like open
    water, which is what separates the Amsterdam canal belt from everything
    else.
    """
    from scipy.spatial import cKDTree

    out = df.copy()
    xy = out[PID].to_numpy(float)
    centre = np.median(xy, axis=0)
    r = np.linalg.norm(xy - centre, axis=1)
    out["radial_px"] = r
    out["radial_frac"] = r / (np.percentile(r, 95) + EPS)

    far = float(np.percentile(r, 99) * 2.0)      # sentinel: "no water anywhere near"
    water = out["sr_MNDWI_max"].to_numpy(float) > 0.0 if "sr_MNDWI_max" in out else None
    if water is not None and water.sum() > 0:
        tree = cKDTree(xy[water])
        d, _ = tree.query(xy, k=1)
        out["dist_to_water_px"] = np.minimum(d, far)
    else:
        out["dist_to_water_px"] = far
    return out


# --------------------------------------------------------------------------- #
# Fallback: start from the organisers' preprocessed_data.pkl
# --------------------------------------------------------------------------- #
def from_official_pkl(path):
    """
    Load the organisers' `preprocessed_data.pkl` into the same shape
    `build_city` produces, so the GBDT notebook can run without the 3 GB of raw
    parquet.

    What you get: the official 60 features, the labels, and the pixel
    coordinates - enough for every classical model, the spatial validation, the
    shift harness and the full few-shot protocol.

    What you do NOT get: the extended features (true surface reflectance,
    construction-event detection, per-decade means, tasseled cap,
    chromaticity, building coverage) and therefore not the physical values the
    verbaliser's word bands need. The pkl's X is already z-scored with Madrid's
    scaler, so the numbers are unitless. Within-city RANKS still work, which is
    the main transfer mechanism - so a ranks-only prompt is possible from this
    file alone, just a thinner one.

    Returns (madrid_df, amsterdam_df).
    """
    import pickle

    with open(path, "rb") as fh:
        d = pickle.load(fh)
    names = list(d["feature_names"])
    out = []
    for city in ("madrid", "amsterdam"):
        X = np.asarray(d[f"X_{city}"], dtype=np.float32)
        y = np.asarray(d[f"y_{city}"], dtype=int)
        ids = np.asarray(d[f"pixel_ids_{city}"])
        df = pd.DataFrame(X, columns=names)
        df.insert(0, "px_key", ids[:, 0])
        df.insert(1, "py_key", ids[:, 1])
        df["age_class"] = y
        df["city"] = city.capitalize()
        out.append(df)
    print(f"loaded {path}")
    for name, df in zip(("Madrid", "Amsterdam"), out):
        print(f"  {name}: {len(df):,} pixels x {len(names)} official features "
              f"| classes {df.age_class.value_counts().sort_index().to_dict()}")
    print("  note: features are Madrid-scaled z-scores, not physical units; "
          "the extended feature set is absent")

    # Guard against loading the wrong file - a truncated download, a synthetic
    # test fixture, or another city's export - and silently reporting numbers
    # from it. The organisers' pixel counts are published in the competitors'
    # guide, so they are a free integrity check.
    import os
    if os.environ.get("ALLOW_SYNTHETIC_PKL") == "1":
        print("  WARNING: integrity check disabled by "
              "ALLOW_SYNTHETIC_PKL - numbers are NOT real")
        return tuple(out)
    expected = {"Madrid": 76_263, "Amsterdam": 25_992}
    problems = []
    for name, df in zip(("Madrid", "Amsterdam"), out):
        if len(df) != expected[name]:
            problems.append(f"{name} has {len(df):,} pixels, expected "
                            f"{expected[name]:,}")
    if len(names) != 60:
        problems.append(f"{len(names)} feature names, expected 60")
    if problems:
        raise ValueError(
            "This does not look like the organisers' preprocessed_data.pkl:\n  "
            + "\n  ".join(problems)
            + f"\nFile: {path}\nIf you meant to use a test fixture, load it "
              "explicitly rather than through this function.")
    print("  integrity check passed (pixel counts match the published figures)")
    return tuple(out)


# --------------------------------------------------------------------------- #
# Driver
# --------------------------------------------------------------------------- #
def build_city(parquet_path, city_hint: str | None = None, verbose: bool = True,
               agg: str = "median", geometry: bool = True) -> pd.DataFrame:
    if verbose:
        print(f"reading {parquet_path}")
    df = pd.read_parquet(parquet_path)
    if "city" not in df.columns and city_hint:
        df["city"] = city_hint
    if verbose:
        print(f"  {len(df):,} rows x {df.shape[1]} cols | "
              f"years {df['year'].min()}-{df['year'].max()}")

    df = assign_age_class(df)
    flat = make_flat_view(df, agg=agg)
    del df
    indicators = period_coverage(flat)

    meta_cols = [c for c in ["city", "coverage", "weighted_mean_year", "age_class"]
                 if c in flat.columns]
    meta = (flat.sort_values(PID + ["year"])
                .groupby(PID, as_index=False)[meta_cols].first())

    pix, years, mats, observed, sat = to_year_matrices(flat)
    del flat

    off = build_official_60(pix, years, mats, indicators)
    ext = build_extended(pix, years, mats, observed, sat)

    feats = off.merge(ext, on=PID, how="left").merge(meta, on=PID, how="left")
    feats = feats[feats["age_class"].notna()].reset_index(drop=True)
    feats["age_class"] = feats["age_class"].astype(int)
    if geometry:
        feats = add_city_geometry(feats)
    if verbose:
        n_ext = ext.shape[1] - len(PID)
        print(f"  -> {len(feats):,} labelled pixels | 60 official + {n_ext} extended features")
    return feats