"""Mapping Lake Mead's shrinking shoreline with Sentinel-2, STAC and NDWI.

Searches the Element 84 Earth Search STAC API for cloud-free Sentinel-2 L2A
scenes over Lake Mead, streams the green / NIR / SWIR / SCL bands from Cloud
Optimised GeoTIFFs at reduced resolution, masks clouds, computes NDWI and
MNDWI and tracks open-water area year by year.
Run with:  python lake_mead_ndwi.py
"""
from collections import defaultdict
from pathlib import Path

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
from matplotlib.colors import ListedColormap
from pyproj import Transformer
from pystac_client import Client
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.vrt import WarpedVRT

HERE = Path(__file__).parent
OUT = HERE / "output"
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 120)

STAC_URL = "https://earth-search.aws.element84.com/v1"
BBOX = (-114.87, 35.98, -114.05, 36.55)      # lon/lat around Lake Mead
TILES = {"11SPA", "11SQA", "11SPV", "11SQV"}  # the four MGRS tiles that cover the lake
YEARS = range(2019, 2026)
RES = 60                                      # analysis resolution in metres
CRS = "EPSG:32611"                            # UTM zone 11N, native CRS of both tiles
CLOUD_CLASSES = [3, 8, 9, 10]                 # SCL: cloud shadow, cloud medium/high, cirrus

# Build a fixed analysis grid so every year lines up pixel-for-pixel.
to_utm = Transformer.from_crs("EPSG:4326", CRS, always_xy=True)
xs, ys = to_utm.transform([BBOX[0], BBOX[2], BBOX[0], BBOX[2]], [BBOX[1], BBOX[1], BBOX[3], BBOX[3]])
x0, x1 = np.floor(min(xs) / RES) * RES, np.ceil(max(xs) / RES) * RES
y0, y1 = np.floor(min(ys) / RES) * RES, np.ceil(max(ys) / RES) * RES
WIDTH, HEIGHT = int((x1 - x0) / RES), int((y1 - y0) / RES)
TRANSFORM = from_origin(x0, y1, RES, RES)


def section(title: str) -> None:
    print(f"\n=== {title} ===")


def tile_id(item) -> str:
    p = item.properties
    return f"{p['mgrs:utm_zone']}{p['mgrs:latitude_band']}{p['mgrs:grid_square']}"


def read_band(href: str, resampling=Resampling.average) -> np.ndarray:
    """Warp a (remote) COG onto the analysis grid. GDAL uses the COG overviews."""
    with rasterio.open(href) as src, WarpedVRT(src, crs=CRS, transform=TRANSFORM, width=WIDTH,
                                               height=HEIGHT, resampling=resampling, nodata=0) as vrt:
        return vrt.read(1)


def scaled(item, band: str) -> np.ndarray:
    """Read a reflectance band and convert digital numbers to surface reflectance.

    Since processing baseline 04.00 ESA adds a +1000 DN offset to L2A data. Earth
    Search v1 re-encodes those COGs with the offset already removed and flags this
    with `earthsearch:boa_offset_applied`; in that case only the scale is needed.
    """
    asset = item.assets[band]
    rb = asset.extra_fields.get("raster:bands", [{}])[0]
    offset = 0.0 if item.properties.get("earthsearch:boa_offset_applied") else rb.get("offset", 0.0)
    dn = read_band(asset.href).astype("float32")
    refl = dn * rb.get("scale", 1.0) + offset
    refl[dn == 0] = np.nan
    return refl


# --- Step 1: search the catalogue -------------------------------------------
catalog = Client.open(STAC_URL)
rows = []
for year in YEARS:
    search = catalog.search(collections=["sentinel-2-l2a"], bbox=BBOX,
                            datetime=f"{year}-07-01/{year}-09-30",
                            query={"eo:cloud_cover": {"lt": 5}})
    by_date = defaultdict(dict)
    for item in search.items():
        tid = tile_id(item)
        if tid not in TILES:
            continue
        # Some dates have been processed twice (e.g. *_0_L2A and a reprocessed *_1_L2A);
        # keep the item with the newest processing baseline.
        prev = by_date[item.datetime.date()].get(tid)
        if prev is None or (item.properties.get("s2:processing_baseline", "")
                            > prev.properties.get("s2:processing_baseline", "")):
            by_date[item.datetime.date()][tid] = item
    complete = {d: v for d, v in by_date.items() if set(v) == TILES}
    if not complete:
        continue

    # Prefer full-swath, clear dates: penalise missing pixels (edge of orbit) and cloud.
    def badness(d):
        return sum(i.properties.get("s2:nodata_pixel_percentage", 0) + i.properties["eo:cloud_cover"]
                   for i in complete[d].values())

    date = min(complete, key=badness)
    items = complete[date]
    rows.append({"year": year, "date": date, "items": items,
                 "cloud_pct": round(np.mean([i.properties["eo:cloud_cover"] for i in items.values()]), 2),
                 "nodata_pct": round(np.mean([i.properties.get("s2:nodata_pixel_percentage", 0)
                                              for i in items.values()]), 2),
                 "baseline": items["11SQA"].properties.get("s2:processing_baseline"),
                 "offset_applied": items["11SQA"].properties.get("earthsearch:boa_offset_applied")})

scenes = pd.DataFrame(rows)
section("selected scenes")
print(scenes.drop(columns="items").to_string(index=False))

example = rows[0]["items"]["11SQA"]
section("example item")
print(example.id)
print("assets:", ", ".join(k for k in example.assets if not k.endswith("-jp2")))
print("green raster:bands:", example.assets["green"].extra_fields["raster:bands"])
print("earthsearch:boa_offset_applied:", example.properties.get("earthsearch:boa_offset_applied"))
latest = rows[-1]["items"]["11SQA"]
print("latest green raster:bands:", latest.assets["green"].extra_fields["raster:bands"])

# --- Step 2: NDWI + cloud mask per year ---------------------------------------
pixel_km2 = RES * RES / 1e6
results, water_masks, rgb = [], {}, {}
for row in rows:
    green = np.full((HEIGHT, WIDTH), np.nan, "float32")
    nir, swir = green.copy(), green.copy()
    scl = np.zeros((HEIGHT, WIDTH), "uint8")
    for item in row["items"].values():  # mosaic the tiles: first valid pixel wins
        g, n, sw = scaled(item, "green"), scaled(item, "nir"), scaled(item, "swir16")
        s = read_band(item.assets["scl"].href, Resampling.nearest)
        fill = np.isnan(green) & ~np.isnan(g)
        green[fill], nir[fill], swir[fill], scl[fill] = g[fill], n[fill], sw[fill], s[fill]

    with np.errstate(invalid="ignore", divide="ignore"):
        ndwi = (green - nir) / (green + nir)      # McFeeters (1996)
        mndwi = (green - swir) / (green + swir)   # Xu (2006): less sensitive to turbidity and glint
    cloudy = np.isin(scl, CLOUD_CLASSES)
    valid = ~np.isnan(ndwi) & ~np.isnan(mndwi) & ~cloudy
    water = (mndwi > 0) & valid
    water_masks[row["year"]] = water
    results.append({"year": row["year"], "date": row["date"],
                    "ndwi_km2": round(((ndwi > 0) & valid).sum() * pixel_km2, 1),
                    "water_km2": round(water.sum() * pixel_km2, 1),
                    "scl_water_km2": round(((scl == 6) & valid).sum() * pixel_km2, 1),
                    "cloud_masked_km2": round(cloudy.sum() * pixel_km2, 1),
                    "valid_pct": round(valid.mean() * 100, 1)})
    print(f"processed {row['date']}: NDWI water = {results[-1]['ndwi_km2']} km², "
          f"MNDWI water = {results[-1]['water_km2']} km²")

    if row["year"] in (rows[0]["year"], 2022):
        tci = np.zeros((3, HEIGHT, WIDTH), "uint8")
        for item in row["items"].values():
            with rasterio.open(item.assets["visual"].href) as src, WarpedVRT(
                    src, crs=CRS, transform=TRANSFORM, width=WIDTH, height=HEIGHT,
                    resampling=Resampling.average, nodata=0) as vrt:
                t = vrt.read()
            empty = tci.sum(axis=0) == 0
            tci[:, empty] = t[:, empty]
        rgb[row["year"]] = np.moveaxis(tci, 0, -1)
        if row["year"] == 2022:
            ndwi_2022, mndwi_2022 = ndwi, mndwi

table = pd.DataFrame(results)
section("water area by year")
print(table.to_string(index=False))
table.to_csv(OUT / "water_area.csv", index=False)

first, last = rows[0]["year"], rows[-1]["year"]
lost = water_masks[first] & ~water_masks[2022]
section("change")
print(f"water in {first} but not 2022: {lost.sum() * pixel_km2:.1f} km²")
print(f"water in 2022 but not {first}: {(water_masks[2022] & ~water_masks[first]).sum() * pixel_km2:.1f} km²")
print(f"change {first} -> {last}: {table.water_km2.iloc[-1] - table.water_km2.iloc[0]:+.1f} km²")

# --- Step 3: figures -------------------------------------------------------------
extent = [x0, x1, y0, y1]
fig, axes = plt.subplots(1, 2, figsize=(11, 4.6))
for ax, year in zip(axes, (first, 2022)):
    ax.imshow(rgb[year], extent=extent)
    ax.set_title(f"Sentinel-2 true colour, {scenes.loc[scenes.year == year, 'date'].iloc[0]}")
    ax.set_xticks([])
    ax.set_yticks([])
fig.savefig(OUT / "01-true-colour.png")
plt.close(fig)

date_2022 = scenes.loc[scenes.year == 2022, "date"].iloc[0]
fig, axes = plt.subplots(1, 2, figsize=(11, 4.4))
for ax, arr, name in zip(axes, (ndwi_2022, mndwi_2022), ("NDWI (green, NIR)", "MNDWI (green, SWIR1)")):
    im = ax.imshow(arr, cmap="RdBu", vmin=-0.8, vmax=0.8, extent=extent)
    ax.set_title(f"{name}, {date_2022}")
    ax.set_xticks([])
    ax.set_yticks([])
fig.colorbar(im, ax=axes, label="Index value (water > 0, blue)", shrink=0.8)
fig.savefig(OUT / "02-ndwi-2022.png")
plt.close(fig)

change = np.zeros((HEIGHT, WIDTH), "uint8")
change[water_masks[first] & water_masks[2022]] = 1
change[lost] = 2
change[water_masks[2022] & ~water_masks[first]] = 3
fig, ax = plt.subplots(figsize=(8, 5.6))
ax.imshow(change, cmap=ListedColormap(["#f4f1ea", "#1d6fa3", "#e8590c", "#2f9e44"]),
          vmin=0, vmax=3, extent=extent, interpolation="nearest")
handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in ("#1d6fa3", "#e8590c", "#2f9e44")]
ax.legend(handles, [f"Water in {first} and 2022", f"Water in {first}, dry in 2022", f"New water in 2022"],
          loc="lower right", fontsize=8)
ax.set_title(f"Shoreline retreat between {first} and 2022 (MNDWI > 0)")
ax.set_xticks([])
ax.set_yticks([])
fig.savefig(OUT / "03-change-map.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8.5, 3.6))
xpos = np.arange(len(table))
for off, col, colour, label in ((-0.27, "water_km2", "#1d6fa3", "MNDWI > 0"),
                                (0.0, "ndwi_km2", "#74c0fc", "NDWI > 0"),
                                (0.27, "scl_water_km2", "#ced4da", "SCL class 6 (water)")):
    bars = ax.bar(xpos + off, table[col], width=0.26, color=colour, label=label)
    if col == "water_km2":
        for b, v in zip(bars, table[col]):
            ax.text(b.get_x() + b.get_width() / 2, v + 2, f"{v:.0f}", ha="center", fontsize=7)
ax.set_xticks(xpos)
ax.set_xticklabels([f"{y}\n{d:%d %b}" for y, d in zip(table.year, table.date)])
ax.set_ylabel("Open water (km²)")
cols = ["water_km2", "ndwi_km2", "scl_water_km2"]
ax.set_ylim(table[cols].min().min() * 0.85, table[cols].max().max() * 1.08)
ax.legend(ncol=3, fontsize=8, loc="upper right")
ax.set_title("Lake Mead open-water area from summer Sentinel-2 scenes")
ax.grid(axis="y", alpha=0.3)
fig.savefig(OUT / "04-water-area-timeseries.png")
plt.close(fig)

# Save the 2022 water mask as a GeoTIFF for use in QGIS.
with rasterio.open(OUT / "water_2022.tif", "w", driver="GTiff", height=HEIGHT, width=WIDTH, count=1,
                   dtype="uint8", crs=CRS, transform=TRANSFORM, compress="deflate") as dst:
    dst.write(water_masks[2022].astype("uint8"), 1)