"""Reproducible Brevard County Landsat surface-temperature change workflow.

The script searches Microsoft Planetary Computer's mirror of USGS Landsat
Collection 2 Level-2 data, reads only the cloud-optimized raster windows needed
for Brevard County, applies pixel-quality masks, aligns two mosaics to a common
60 m grid, and calculates a paired-pixel temperature difference.

This is a two-date comparison, not a climate-trend analysis.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import planetary_computer
import rasterio
import requests
from matplotlib.colors import TwoSlopeNorm
from pystac_client import Client
from rasterio.features import geometry_mask
from rasterio.transform import from_origin
from rasterio.windows import Window, from_bounds
from rasterio.warp import Resampling, reproject, transform_geom
from shapely.geometry import shape


STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
CENSUS_URL = (
    "https://tigerweb.geo.census.gov/arcgis/rest/services/"
    "TIGERweb/State_County/MapServer/1/query"
)
COLLECTION = "landsat-c2-l2"
TARGET_CRS = "EPSG:32617"
TARGET_RESOLUTION = 60
ST_SCALE = 0.00341802
ST_OFFSET_C = 149.0 - 273.15

SCENES = {
    "2009": {
        "date": "2009-02-13 to 2009-02-20",
        "platform": "Landsat 7",
        "ids": [
            "LE07_L2SP_015040_20090213_02_T1",
            "LE07_L2SP_015041_20090213_02_T1",
            "LE07_L2SP_016040_20090220_02_T1",
            "LE07_L2SP_016041_20090220_02_T1",
        ],
    },
    "2025": {
        "date": "2025-02-09",
        "platform": "Landsat 9",
        "ids": [
            "LC09_L2SP_015040_20250209_02_T1",
            "LC09_L2SP_015041_20250209_02_T1",
        ],
    },
}

ROOT = Path(__file__).resolve().parents[2]
MAP_OUTPUT = ROOT / "public" / "projects" / "python-raster-change.png"
SUMMARY_OUTPUT = ROOT / "public" / "downloads" / "landsat_change_summary.json"


def brevard_geometry() -> dict:
    """Retrieve Brevard County's current TIGERweb boundary as GeoJSON."""
    response = requests.get(
        CENSUS_URL,
        params={
            "where": "STATE='12' AND COUNTY='009'",
            "outFields": "NAME,STATE,COUNTY",
            "returnGeometry": "true",
            "f": "geojson",
        },
        timeout=90,
    )
    response.raise_for_status()
    return response.json()["features"][0]["geometry"]


def target_grid(geometry_wgs84: dict) -> tuple[dict, rasterio.Affine, int, int]:
    """Create a county-wide 60 m grid snapped to whole resolution units."""
    geometry = transform_geom("EPSG:4326", TARGET_CRS, geometry_wgs84)
    min_x, min_y, max_x, max_y = shape(geometry).bounds
    left = math.floor(min_x / TARGET_RESOLUTION) * TARGET_RESOLUTION
    bottom = math.floor(min_y / TARGET_RESOLUTION) * TARGET_RESOLUTION
    right = math.ceil(max_x / TARGET_RESOLUTION) * TARGET_RESOLUTION
    top = math.ceil(max_y / TARGET_RESOLUTION) * TARGET_RESOLUTION
    width = round((right - left) / TARGET_RESOLUTION)
    height = round((top - bottom) / TARGET_RESOLUTION)
    return geometry, from_origin(left, top, TARGET_RESOLUTION, TARGET_RESOLUTION), width, height


def get_items(scene_ids: list[str]) -> list:
    """Find and sign the requested Landsat scene items."""
    catalog = Client.open(STAC_URL, modifier=planetary_computer.sign_inplace)
    items = []
    for scene_id in scene_ids:
        search = catalog.search(collections=[COLLECTION], ids=[scene_id])
        matches = list(search.items())
        if len(matches) != 1:
            raise RuntimeError(f"Expected one STAC item for {scene_id}; found {len(matches)}")
        items.append(matches[0])
    return items


def clipped_window(dataset: rasterio.DatasetReader, target_bounds: tuple[float, ...]) -> Window:
    """Return the source window intersecting the county grid bounds."""
    requested = from_bounds(*target_bounds, transform=dataset.transform).round_offsets().round_lengths()
    full = Window(0, 0, dataset.width, dataset.height)
    return requested.intersection(full)


def read_asset_window(href: str, bounds: tuple[float, ...]) -> tuple[np.ndarray, rasterio.Affine, object]:
    """Read only the COG window overlapping Brevard County."""
    with rasterio.open(href) as source:
        window = clipped_window(source, bounds)
        return source.read(1, window=window), source.window_transform(window), source.crs


def clear_land_mask(qa_pixel: np.ndarray, qa_radsat: np.ndarray, temperature_dn: np.ndarray) -> np.ndarray:
    """Mask fill, dilated cloud, cloud, shadow, snow, water, and saturation/dropouts."""
    rejected_bits = sum(1 << bit for bit in (0, 1, 3, 4, 5, 7))
    return (
        ((qa_pixel & rejected_bits) == 0)
        & (qa_radsat == 0)
        & (temperature_dn > 0)
    )


def build_mosaic(items: list, transform: rasterio.Affine, width: int, height: int) -> np.ndarray:
    """Quality-mask, reproject, and average overlapping row 40/41 observations."""
    sums = np.zeros((height, width), dtype="float64")
    counts = np.zeros((height, width), dtype="uint8")
    bounds = rasterio.transform.array_bounds(height, width, transform)

    for item in items:
        thermal_key = "lwir11" if "lwir11" in item.assets else "lwir"
        temperature_dn, source_transform, source_crs = read_asset_window(item.assets[thermal_key].href, bounds)
        qa_pixel, qa_transform, _ = read_asset_window(item.assets["qa_pixel"].href, bounds)
        qa_radsat, radsat_transform, _ = read_asset_window(item.assets["qa_radsat"].href, bounds)
        if qa_transform != source_transform or radsat_transform != source_transform:
            raise RuntimeError(f"Quality bands are not aligned for {item.id}")

        valid = clear_land_mask(qa_pixel, qa_radsat, temperature_dn)
        temperature_c = temperature_dn.astype("float32") * ST_SCALE + ST_OFFSET_C
        source = np.where(valid, temperature_c, -9999.0).astype("float32")
        aligned = np.full((height, width), -9999.0, dtype="float32")
        reproject(
            source=source,
            destination=aligned,
            src_transform=source_transform,
            src_crs=source_crs,
            src_nodata=-9999.0,
            dst_transform=transform,
            dst_crs=TARGET_CRS,
            dst_nodata=-9999.0,
            resampling=Resampling.bilinear,
        )
        usable = aligned > -1000
        sums[usable] += aligned[usable]
        counts[usable] += 1

    return np.divide(sums, counts, out=np.full_like(sums, np.nan), where=counts > 0).astype("float32")


def outline_coordinates(geometry_projected: dict) -> list[tuple[np.ndarray, np.ndarray]]:
    """Extract exterior rings for plotting Polygon or MultiPolygon geometry."""
    polygons = geometry_projected["coordinates"]
    if geometry_projected["type"] == "Polygon":
        polygons = [polygons]
    lines = []
    for polygon in polygons:
        exterior = np.asarray(polygon[0])
        lines.append((exterior[:, 0], exterior[:, 1]))
    return lines


def write_map(
    difference: np.ndarray,
    transform: rasterio.Affine,
    geometry_projected: dict,
    stats: dict,
) -> None:
    """Render the portfolio preview and project-page analytical figure."""
    plt.rcParams.update({"font.family": "DejaVu Sans"})
    fig = plt.figure(figsize=(16, 10), facecolor="#f8f6ef")
    grid = fig.add_gridspec(1, 2, width_ratios=[1.72, 0.72], wspace=0.035)
    ax = fig.add_subplot(grid[0, 0], facecolor="#e7ece9")
    panel = fig.add_subplot(grid[0, 1], facecolor="#173f52")

    finite = difference[np.isfinite(difference)]
    limit = max(4.0, float(np.nanpercentile(np.abs(finite), 95)))
    left, bottom, right, top = rasterio.transform.array_bounds(*difference.shape, transform)
    image = ax.imshow(
        difference,
        extent=(left, right, bottom, top),
        origin="upper",
        cmap="RdBu_r",
        norm=TwoSlopeNorm(vmin=-limit, vcenter=0, vmax=limit),
    )
    for x, y in outline_coordinates(geometry_projected):
        ax.plot(x, y, color="#17242c", linewidth=1.2)
    ax.set_title("Brevard County surface-temperature difference", loc="left", fontsize=22, fontweight="bold", pad=18, color="#17242c")
    ax.text(0, 1.012, "LANDSAT 9 (2025) MINUS LANDSAT 7 (2009) · QA-MASKED PAIRED PIXELS", transform=ax.transAxes, fontsize=9.5, fontweight="bold", color="#2c6f84")
    ax.set_axis_off()
    colorbar = fig.colorbar(image, ax=ax, orientation="horizontal", fraction=0.035, pad=0.035, aspect=35)
    colorbar.set_label("Surface-temperature difference (°C) · blue = cooler, red = warmer", fontsize=10, fontweight="bold")

    panel.set_xticks([])
    panel.set_yticks([])
    for spine in panel.spines.values():
        spine.set_visible(False)
    panel.text(0.08, 0.93, "PYTHON WORKFLOW", color="#9cc9cd", fontsize=10, fontweight="bold", transform=panel.transAxes)
    mean_display = 0.0 if abs(stats["mean_difference_c"]) < 0.05 else stats["mean_difference_c"]
    panel.text(0.08, 0.855, f"{mean_display:+.1f}°C", color="white", fontsize=39, fontweight="bold", transform=panel.transAxes)
    panel.text(0.08, 0.81, "mean paired-pixel difference", color="#c8d6da", fontsize=10, transform=panel.transAxes)
    panel.text(0.08, 0.715, f"{stats['median_difference_c']:+.1f}°C", color="white", fontsize=28, fontweight="bold", transform=panel.transAxes)
    panel.text(0.08, 0.68, "median difference", color="#c8d6da", fontsize=10, transform=panel.transAxes)
    panel.text(0.08, 0.59, f"{stats['paired_pixels']:,}", color="white", fontsize=27, fontweight="bold", transform=panel.transAxes)
    panel.text(0.08, 0.555, "clear land pixels compared", color="#c8d6da", fontsize=10, transform=panel.transAxes)
    panel.plot([0.08, 0.92], [0.49, 0.49], color="#557181", linewidth=1, transform=panel.transAxes)
    steps = [
        "01  Search official\n      Landsat metadata",
        "02  Read cloud-optimized\n      raster windows",
        "03  Mask clouds, water,\n      gaps & saturation",
        "04  Align mosaics to a\n      common 60 m grid",
        "05  Subtract, summarize\n      & visualize",
    ]
    for index, step in enumerate(steps):
        panel.text(0.08, 0.435 - index * 0.075, step, color="white", fontsize=9.6, fontweight="bold", linespacing=1.3, transform=panel.transAxes)
    panel.text(0.08, 0.07, "Two-date comparison—not a climate trend.", color="#efb28d", fontsize=9.5, fontweight="bold", transform=panel.transAxes)
    panel.text(0.08, 0.037, "Acquisition weather and sensor differences remain.", color="#c8d6da", fontsize=8.6, transform=panel.transAxes)
    fig.savefig(MAP_OUTPUT, dpi=150, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)


def main() -> None:
    geometry_wgs84 = brevard_geometry()
    geometry_projected, transform, width, height = target_grid(geometry_wgs84)
    county = geometry_mask([geometry_projected], (height, width), transform, invert=True)

    mosaics = {}
    for year, scene in SCENES.items():
        mosaics[year] = build_mosaic(get_items(scene["ids"]), transform, width, height)

    difference = mosaics["2025"] - mosaics["2009"]
    difference[~county] = np.nan
    values = difference[np.isfinite(difference)]
    if values.size == 0:
        raise RuntimeError("No paired clear land pixels remained after quality masking")

    stats = {
        "study_area": "Brevard County, Florida",
        "earlier_scene": f"{SCENES['2009']['platform']} · {SCENES['2009']['date']}",
        "later_scene": f"{SCENES['2025']['platform']} · {SCENES['2025']['date']}",
        "grid_resolution_m": TARGET_RESOLUTION,
        "paired_pixels": int(values.size),
        "mean_difference_c": round(float(np.mean(values)), 2),
        "median_difference_c": round(float(np.median(values)), 2),
        "percent_above_plus_2_c": round(float(np.mean(values > 2) * 100), 1),
        "percent_below_minus_2_c": round(float(np.mean(values < -2) * 100), 1),
        "difference_p10_c": round(float(np.percentile(values, 10)), 2),
        "difference_p90_c": round(float(np.percentile(values, 90)), 2),
        "interpretation": "Two-date surface-temperature comparison; not a climate-trend estimate.",
        "source": "USGS Landsat Collection 2 Level-2, accessed through Microsoft Planetary Computer",
    }
    MAP_OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    SUMMARY_OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    write_map(difference, transform, geometry_projected, stats)
    SUMMARY_OUTPUT.write_text(json.dumps(stats, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(stats, indent=2))


if __name__ == "__main__":
    main()
