# %% [markdown]
# # 1 - Features: from 40 years of Landsat to one row per building pixel
#
# This notebook turns the two raw parquet files into two feature tables.
#
# It does **two** things at once:
#
# 1. **Reproduces the organisers' 60 features exactly** (Notebook 3), so every
#    number we report stays comparable with the published Random-Forest
#    baseline of 0.618 +- 0.004 macro F1 on Madrid.
# 2. **Adds an extended set** on top: proper surface reflectance instead of raw
#    digital numbers, correct spectral indices, construction-event detection,
#    per-decade means, trends, and label-free neighbourhood context.
#
# Nothing here touches `y_amsterdam`. Amsterdam's *features* are used (for
# per-city statistics), its *labels* are not.
#
# Runs on CPU in a few minutes - no GPU needed. Do this once, locally, and the
# GPU notebooks just load the parquet files it writes.

# %%
import sys
from pathlib import Path

# Make src/ importable whether we are local, on Kaggle, or on Colab.
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 features as F
import retrieval as R

pd.set_option("display.width", 200)
pd.set_option("display.max_columns", 60)
print("platform:", C.PLATFORM)
print("looking for raw parquet under:", C.RAW_DIR)

# %% [markdown]
# ## 1.1 Find the raw files
#
# Expected: `madrid_train.parquet` (about 3.17 M rows) and
# `amsterdam_data.parquet` (about 1.06 M rows). On Kaggle they will be inside
# `/kaggle/input/<dataset-name>/`; the search below handles that.

# %%
def find_parquet(pattern: str) -> Path:
    roots = [C.RAW_DIR, C.DATA_DIR, C.PROJECT_ROOT, Path("/kaggle/input")]
    for root in roots:
        if not root.exists():
            continue
        hits = sorted(root.rglob(pattern))
        if hits:
            return hits[0]
    raise FileNotFoundError(
        f"{pattern} not found. Put the two parquet files in {C.RAW_DIR} "
        "(local) or attach them as a Kaggle dataset.")

MADRID_PARQUET = find_parquet("madrid_train.parquet")
AMS_PARQUET = find_parquet("amsterdam_data.parquet")
print("Madrid   :", MADRID_PARQUET)
print("Amsterdam:", AMS_PARQUET)

# %% [markdown]
# ## 1.2 What is actually in the raw file
#
# One row per *(pixel, year)*. Up to three summer observations per year are
# stored side by side as `Blue_1/Blue_2/Blue_3`, each with a `qa_valid_N`
# cloud flag. Band values are raw Landsat Collection-2 integers:
#
# $$\text{surface reflectance} = \text{DN} \times 2.75\times10^{-5} - 0.2$$

# %%
peek = pd.read_parquet(MADRID_PARQUET)
print(f"{len(peek):,} rows x {peek.shape[1]} columns")
print("\ncolumns:")
print(list(peek.columns))
print("\nyears:", peek["year"].min(), "-", peek["year"].max(),
      f"({peek['year'].nunique()} distinct)")
print("unique pixels:", f"{peek.groupby(C.PID).ngroups:,}")
print("\nfirst row:")
print(peek.iloc[0].to_string())

# %%
# Reflectance sanity check on one band, and the QA-flag hit rate.
dn = peek["Red_1"].dropna()
sr = dn * C.SR_SCALE + C.SR_OFFSET
print("Red_1 digital numbers :", dn.quantile([0, .01, .5, .99, 1]).round(0).to_dict())
print("Red_1 as reflectance  :", sr.quantile([0, .01, .5, .99, 1]).round(3).to_dict())
for s in (1, 2, 3):
    col = f"qa_valid_{s}"
    if col in peek.columns:
        print(f"{col}: {peek[col].mean():.1%} valid")
del peek

# %% [markdown]
# ## 1.3 Build both feature sets
#
# The heavy lifting lives in `src/features.py`. Per city it:
#
# * picks the first QA-valid observation per band per year, drops haze
#   (`Blue > 15000`);
# * records `has_early_data` / `has_late_data` on the *observed* years;
# * fills the annual series (linear inside, nearest at the edges);
# * pivots to one `(pixels x years)` matrix per band, then computes every
#   statistic with numpy - a whole city in about a minute.

# %%
madrid = F.build_city(MADRID_PARQUET, city_hint="Madrid")

# %%
amsterdam = F.build_city(AMS_PARQUET, city_hint="Amsterdam")

# %% [markdown]
# ## 1.4 Does this match the organisers' own output?
#
# If `preprocessed_data.pkl` from the organisers' Notebook 3 is around, compare
# the 60 official features against it. Their arrays are standardised with
# Madrid's `StandardScaler`, so we compare **correlation per feature**, which is
# invariant to that. Anything below 0.999 means our pipeline has drifted.

# %%
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:
    print("preprocessed_data.pkl not found - skipping the cross-check.")
else:
    import pickle
    with open(pkl, "rb") as fh:
        ref = pickle.load(fh)
    print("reference file:", pkl)
    ours = madrid.set_index(["px_key", "py_key"])
    ref_ids = pd.MultiIndex.from_arrays(
        [ref["pixel_ids_madrid"][:, 0], ref["pixel_ids_madrid"][:, 1]],
        names=["px_key", "py_key"])
    common = ref_ids.intersection(ours.index)
    print(f"pixels: ours {len(ours):,}, theirs {len(ref_ids):,}, "
          f"in common {len(common):,}")
    ref_df = pd.DataFrame(ref["X_madrid"], index=ref_ids,
                          columns=ref["feature_names"]).loc[common]
    our_df = ours.loc[common, ref["feature_names"]]
    corr = {c: float(np.corrcoef(ref_df[c], our_df[c])[0, 1])
            for c in ref["feature_names"]
            if ref_df[c].std() > 0 and our_df[c].std() > 0}
    bad = {k: v for k, v in corr.items() if v < 0.999}
    print(f"features compared: {len(corr)}  |  min correlation: "
          f"{min(corr.values()):.5f}")
    print("dead features in their file (zero variance):",
          [c for c in ref["feature_names"] if ref_df[c].std() == 0])
    print("features below 0.999:", bad if bad else "none - pipeline matches")
    # Labels must agree exactly.
    y_ref = pd.Series(ref["y_madrid"], index=ref_ids).loc[common]
    print("label agreement:",
          f"{(y_ref.to_numpy() == our_df.index.map(ours['age_class']).to_numpy()).mean():.4%}")

# %% [markdown]
# ## 1.5 Class balance, and the domain gap
#
# Two facts drive every later design decision:
#
# * the classes are unbalanced, differently in each city;
# * the same feature has a different *level* in each city, by much more than
#   the gap between classes - which is why a Madrid model collapses to ~0.34 on
#   Amsterdam.

# %%
for name, df in [("Madrid", madrid), ("Amsterdam", amsterdam)]:
    n = len(df)
    counts = df["age_class"].value_counts().sort_index()
    print(f"\n{name}: {n:,} labelled pixels")
    for c in C.CLASS_LABELS[name]:
        k = int(counts.get(c, 0))
        print(f"  class {c} ({C.CLASS_LABELS[name][c]:>11}): {k:6,}  {k / n:6.1%}")

# %%
# How far apart are the two cities, feature by feature?
# Expressed in Madrid standard deviations, so it is directly comparable.
key = ["sr_brightness_mean", "sr_Red_mean", "sr_NIR_mean", "sr_SWIR1_mean",
       "sr_NDVI_mean", "sr_NDBI_mean", "sr_UI_mean", "sr_MNDWI_mean",
       "sr_BSI_mean", "sr_brightness_std", "sr_brightness_late_minus_early",
       "step_brightness_mag", "step_brightness_njumps", "coverage"]
key = [k for k in key if k in madrid.columns and k in amsterdam.columns]
gap = pd.DataFrame({
    "madrid_mean": [madrid[k].mean() for k in key],
    "amsterdam_mean": [amsterdam[k].mean() for k in key],
}, index=key)
gap["madrid_sd"] = [madrid[k].std() for k in key]
gap["gap_in_madrid_sd"] = (gap.amsterdam_mean - gap.madrid_mean) / gap.madrid_sd
print(gap.round(3).sort_values("gap_in_madrid_sd").to_string())

# %% [markdown]
# A feature whose `gap_in_madrid_sd` is near zero is one a Madrid-trained model
# can still read in Amsterdam. A feature at +/-2 is telling the model which
# *city* it is looking at, not how old the building is.

# %%
# Within-class spread versus between-city gap, for the top offenders.
fig, axes = plt.subplots(2, 4, figsize=(18, 7))
for ax, k in zip(axes.ravel(), key[:8]):
    for name, df, col in [("Madrid", madrid, "tab:red"),
                          ("Amsterdam", amsterdam, "tab:blue")]:
        ax.hist(df[k], bins=60, density=True, alpha=0.45, label=name, color=col)
    ax.set_title(k, fontsize=9)
    ax.tick_params(labelsize=7)
axes.ravel()[0].legend(fontsize=8)
plt.suptitle("The same feature in the two cities", fontweight="bold")
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 1.6 Does the construction-event detector actually work?
#
# A genuine test of the extended features. For classes 3 and 4 the building went
# up *during* the satellite record, so `step_brightness_year` should land near
# `weighted_mean_year`. For classes 1 and 2 it was already standing in 1984, so
# the detector should find nothing meaningful. If that pattern does not show up,
# the detector is noise and should be dropped.

# %%
for name, df in [("Madrid", madrid), ("Amsterdam", amsterdam)]:
    print(f"\n{name}")
    for sig in ("brightness", "SWIR1", "NDVI"):
        sub = df[df.age_class.isin([3, 4])]
        err = (sub[f"step_{sig}_year"] - sub["weighted_mean_year"]).abs()
        r = np.corrcoef(sub[f"step_{sig}_year"], sub["weighted_mean_year"])[0, 1]
        print(f"  step_{sig}_year vs true year (classes 3+4): r={r:+.3f}  "
              f"median |error| = {err.median():.1f} years  "
              f"within 5 y: {(err <= 5).mean():.1%}")
    print("  mean |step magnitude| by class:",
          df.groupby("age_class")["step_brightness_mag"].apply(
              lambda s: round(float(s.abs().mean()), 4)).to_dict())

# %%
fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))
for ax, (name, df) in zip(axes, [("Madrid", madrid), ("Amsterdam", amsterdam)]):
    sub = df[df.age_class.isin([3, 4])].sample(min(4000, len(df)), random_state=0)
    ax.scatter(sub["weighted_mean_year"], sub["step_brightness_year"], s=3,
               alpha=0.25)
    lims = [1984, 2024]
    ax.plot(lims, lims, "r--", lw=1)
    ax.set(xlabel="true construction year", ylabel="detected brightness step year",
           title=f"{name}: classes 3 and 4", xlim=lims, ylim=lims)
plt.tight_layout()
plt.show()

# %% [markdown]
# ## 1.7 Neighbourhood context - free information, no labels
#
# Buildings go up a street at a time, so a pixel's neighbours carry real signal
# about its era. These features average a few key quantities over the 8 nearest
# pixels (within 10 pixel units = 300 m) and also store the *difference* between
# the pixel and its surroundings. No labels are involved, so this is available
# in Amsterdam at n = 0.

# %%
madrid = R.add_neighbourhood_features(madrid)
amsterdam = R.add_neighbourhood_features(amsterdam)
nbr_cols = [c for c in madrid.columns if c.startswith("nbr_")]
print(f"added {len(nbr_cols)} neighbourhood features")
print(madrid[["age_class"] + nbr_cols[:6]].groupby("age_class").mean().round(4).to_string())

# %% [markdown]
# ## 1.8 Save
#
# One parquet per city. Every later notebook starts from these.

# %%
C.PROC_DIR.mkdir(parents=True, exist_ok=True)
madrid.to_parquet(C.PROC_DIR / "madrid_features.parquet", index=False)
amsterdam.to_parquet(C.PROC_DIR / "amsterdam_features.parquet", index=False)
for n, df in [("madrid", madrid), ("amsterdam", amsterdam)]:
    p = C.PROC_DIR / f"{n}_features.parquet"
    print(f"{p}  {p.stat().st_size / 1e6:.1f} MB  {df.shape}")
print("\nfeature groups:",
      f"{len(C.FEATURE_COLS_60)} official,",
      f"{sum(1 for c in madrid.columns if c.startswith('sr_'))} reflectance-based,",
      f"{sum(1 for c in madrid.columns if c.startswith('step_'))} change-event,",
      f"{len(nbr_cols)} neighbourhood")