Remote sensing data

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import rasterio
from pyproj import Transformer
from rasterio.windows import from_bounds
from rasterio.windows import transform as window_transform


def bbox_to_dataset_crs(bbox_wgs84, dst_crs):
    """Transform a WGS84 bounding box to a raster's native CRS.

    Rasterio windows must be given in the source file's own CRS, not WGS84.

    Args:
        bbox_wgs84: Bounding box in WGS84 [min_lon, min_lat, max_lon, max_lat].
        dst_crs: Target CRS (e.g., src.crs from rasterio.open()).

    Returns:
        Bounding box in the target CRS [min_x, min_y, max_x, max_y].
    """
    transformer = Transformer.from_crs("EPSG:4326", dst_crs, always_xy=True)
    min_lon, min_lat, max_lon, max_lat = bbox_wgs84
    xs, ys = transformer.transform([min_lon, max_lon], [min_lat, max_lat])
    return min(xs), min(ys), max(xs), max(ys)


def ycbcr_to_rgb(stack):
    """Convert a 3-band YCbCr image (ITU-R BT.601) to RGB uint8.

    Some imagery (e.g. JPEG-compressed GeoTIFFs, as used by SWISSIMAGE) stores
    colour as YCbCr rather than RGB, and must be converted before it can be
    written or plotted as a normal RGB image.

    Args:
        stack: Array of shape (3, height, width) with YCbCr bands.

    Returns:
        RGB array of shape (3, height, width) as uint8.
    """
    y, cb, cr = stack.astype("float32")
    r = y + 1.402 * (cr - 128.0)
    g = y - 0.344136 * (cb - 128.0) - 0.714136 * (cr - 128.0)
    b = y + 1.772 * (cb - 128.0)
    return np.clip(np.stack([r, g, b]), 0, 255).astype("uint8")


def save_single_asset_subset(
    item, asset_key, bbox_wgs84, output_path, band_indexes=(1, 2, 3)
):
    """Read the AOI window from one raster asset and save it as a GeoTIFF.

    Output pixel dimensions follow the AOI at the source's native resolution,
    so a landscape bbox produces a landscape image (no separate width/height
    arguments needed).

    Args:
        item: A STAC Item object.
        asset_key: String key of the asset to read (e.g., 'B02', 'red', 'image').
        bbox_wgs84: Bounding box in WGS84 [min_lon, min_lat, max_lon, max_lat].
        output_path: File path for the output GeoTIFF.
        band_indexes: Tuple of band indices to read (1-indexed; default (1,2,3) for RGB).

    Returns:
        Path to the written GeoTIFF file.
    """
    href = item.assets[asset_key].href

    with rasterio.open(href) as src:
        print("Source CRS:", src.crs)
        bounds = bbox_to_dataset_crs(bbox_wgs84, src.crs)
        window = from_bounds(*bounds, transform=src.transform)

        stack = src.read(
            list(band_indexes), window=window, boundless=True, fill_value=0
        )
        print("Read shape, dtype:", stack.shape, stack.dtype)

        profile = src.profile.copy()
        profile.update(
            driver="GTiff",
            height=stack.shape[1],
            width=stack.shape[2],
            count=len(band_indexes),
            transform=window_transform(window, src.transform),
            compress="deflate",
        )
        profile.pop("photometric", None)

        if (
            str(src.profile.get("photometric", "")).lower() == "ycbcr"
            and stack.shape[0] == 3
        ):
            print("Converting YCbCr -> RGB")
            stack = ycbcr_to_rgb(stack)
        profile.update(dtype=str(stack.dtype))

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(output_path, "w", **profile) as dst:
        dst.write(stack)

    size_mb = output_path.stat().st_size / 1_000_000
    print(f"Saved: {output_path} ({size_mb:.2f} MB)")
    return output_path


def save_multiband_subset(item, asset_keys, bbox_wgs84, output_path):
    """Read the AOI window from several single-band assets and save as one GeoTIFF.

    Each asset key is read separately (one band each) and stacked in the given
    order, e.g. asset_keys=['B04', 'B03', 'B02'] for a red-green-blue stack.
    Output pixel dimensions follow the AOI at the source's native resolution.

    Args:
        item: A STAC Item object.
        asset_keys: Ordered list of asset keys, one per output band.
        bbox_wgs84: Bounding box in WGS84 [min_lon, min_lat, max_lon, max_lat].
        output_path: File path for the output GeoTIFF.

    Returns:
        Path to the written GeoTIFF file.
    """
    bands = []
    profile = None

    for key in asset_keys:
        with rasterio.open(item.assets[key].href) as src:
            bounds = bbox_to_dataset_crs(bbox_wgs84, src.crs)
            window = from_bounds(*bounds, transform=src.transform)
            band = src.read(1, window=window, boundless=True, fill_value=0)
            bands.append(band)

            if profile is None:
                profile = src.profile.copy()
                profile.update(
                    driver="GTiff",
                    height=band.shape[0],
                    width=band.shape[1],
                    count=len(asset_keys),
                    transform=window_transform(window, src.transform),
                    compress="deflate",
                )
                profile.pop("photometric", None)

    stack = np.stack(bands)
    profile.update(dtype=str(stack.dtype))

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(output_path, "w", **profile) as dst:
        dst.write(stack)

    size_mb = output_path.stat().st_size / 1_000_000
    print(f"Saved: {output_path} ({size_mb:.2f} MB)")
    return output_path


def stretch_rgb(rgb):
    """Apply a percentile stretch to an RGB image for better visual contrast.

    Raw satellite data often spans a wide dynamic range. A 2-98 percentile stretch clips
    extreme values and maps the middle 96% to the display range, improving visibility of
    detail without losing important information to saturation.

    Args:
        rgb: Array of shape (3, height, width) or (height, width, 3).

    Returns:
        Stretched array normalized to [0, 1].
    """
    rgb = np.moveaxis(rgb, 0, -1).astype("float32")
    p_low, p_high = np.nanpercentile(rgb, (2, 98))
    if p_high == p_low:
        return np.zeros_like(rgb)
    rgb = (rgb - p_low) / (p_high - p_low)
    return np.clip(rgb, 0, 1)


def plot_rgb(path, title):
    """Load a GeoTIFF and display it as an RGB image with percentile stretch.

    Args:
        path: File path to a 3-band GeoTIFF.
        title: Title for the plot.
    """
    with rasterio.open(path) as src:
        rgb = src.read([1, 2, 3])
    plt.figure(figsize=(7, 3))
    plt.imshow(stretch_rgb(rgb))
    plt.title(title)
    plt.axis("off")
    plt.show()

from pystac_client import Client
import planetary_computer as pc

pc_catalog = Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

rome_bbox = [12.42, 41.88, 12.52, 41.91]  # WGS84 lon/lat around Rome

search = pc_catalog.search(
    collections=["sentinel-2-l2a"],
    bbox=rome_bbox,
    datetime="2025-06-01/2025-08-31",
    query={"eo:cloud_cover": {"lt": 20}},
    max_items=10,
)

items = list(search.items())
print(f"Found {len(items)} Sentinel-2 item(s)")

items_sorted = sorted(
    items,
    key=lambda item: item.properties.get("eo:cloud_cover", 999),
)

s2_item = pc.sign(items_sorted[0])

print("Selected item:", s2_item.id)
print("Date:", s2_item.datetime)
print("Cloud cover:", round(s2_item.properties.get("eo:cloud_cover"), 2))

print("\nUseful assets:")
for key in ["B02", "B03", "B04", "B08", "SCL", "visual"]:
    if key in s2_item.assets:
        print(key, "-", s2_item.assets[key].title)

s2_subset = save_multiband_subset(
    item=s2_item,
    asset_keys=["B04", "B03", "B02"],  # red, green, blue
    bbox_wgs84=rome_bbox,
    output_path="data/raw/sentinel2/s2_rome_rgb_subset.tif",
)

plot_rgb(s2_subset, "Sentinel-2 subset: Rome")

everest_bbox = [86.80, 27.90, 87.01, 27.98]  # WGS84 lon/lat around Mount Everest

search = pc_catalog.search(
    collections=["landsat-c2-l2"],
    bbox=everest_bbox,
    datetime="2026-07-01/2026-09-30",
    query={"eo:cloud_cover": {"lt": 30}},
    max_items=10,
)

items = list(search.items())
print(f"Found {len(items)} Landsat item(s)")

items_sorted = sorted(
    items,
    key=lambda item: item.properties.get("eo:cloud_cover", 999),
)

landsat_item = pc.sign(items_sorted[0])

print("Selected item:", landsat_item.id)
print("Date:", landsat_item.datetime)
print("Cloud cover:", landsat_item.properties.get("eo:cloud_cover"))

print("\nFirst asset keys:")
print(list(landsat_item.assets.keys())[:20])

landsat_subset = save_multiband_subset(
    item=landsat_item,
    asset_keys=["swir22", "nir08", "red"],
    bbox_wgs84=everest_bbox,
    output_path="data/raw/landsat/landsat_everest_rgb_subset.tif",
)

plot_rgb(landsat_subset, "Landsat subset: Mount Everest region")

nyc_bbox = [-74.050, 40.688, -74.039, 40.691]  # WGS84 lon/lat around Statue of Liberty

search = pc_catalog.search(
    collections=["naip"],
    bbox=nyc_bbox,
    datetime="2021-01-01/2024-12-31",
    max_items=10,
)

items = list(search.items())
print(f"Found {len(items)} NAIP item(s)")

naip_item = pc.sign(items[0])

print("Selected item:", naip_item.id)
print("Date:", naip_item.datetime)
print("\nAsset keys:")
print(list(naip_item.assets.keys()))

naip_asset_key = "image"

if naip_asset_key not in naip_item.assets:
    raise KeyError(f"Check available NAIP assets: {list(naip_item.assets.keys())}")

naip_subset = save_single_asset_subset(
    item=naip_item,
    asset_key=naip_asset_key,
    bbox_wgs84=nyc_bbox,
    output_path="data/raw/naip/naip_statue_of_liberty_rgb_subset.tif",
    band_indexes=(1, 2, 3),  # red, green, blue
)

plot_rgb(naip_subset, "NAIP subset: Statue of Liberty, New York")

from pystac_client import Client

vantor_catalog = Client.open(
    "https://vantor-opendata.s3.amazonaws.com/events/catalog.json"
)

collections = list(vantor_catalog.get_collections())
print(f"Found {len(collections)} Vantor event collection(s)")

for collection in collections[:10]:
    print(collection.id, "-", collection.title)

matching = [
    collection
    for collection in collections
    if "Bordeaux" in collection.id or "Bordeaux" in (collection.title or "")
]

if not matching:
    raise ValueError(
        "No Bordeaux collection found. Choose another event from the printed list."
    )

vantor_collection = matching[0]

print("Selected collection:", vantor_collection.id)
print("Title:", vantor_collection.title)
print("License:", vantor_collection.license)
print("Temporal extent:", vantor_collection.extent.temporal.intervals)
print("Spatial extent:", vantor_collection.extent.spatial.bboxes)

vantor_items = list(vantor_collection.get_items())
print(f"Found {len(vantor_items)} item(s)")

vantor_item = vantor_items[0]

print("Selected item:", vantor_item.id)
print("Date:", vantor_item.datetime)
print("Item bbox:", vantor_item.bbox)

print("\nAssets:")
for key, asset in vantor_item.assets.items():
    print(key, "-", asset.title, "-", asset.media_type)

vantor_bbox = [-1.0846, 44.8698, -1.0666, 44.8748]

print(vantor_bbox)

image_asset_keys = [
    key
    for key, asset in vantor_item.assets.items()
    if "tif" in str(asset.href).lower() or "tiff" in str(asset.media_type).lower()
]

if not image_asset_keys:
    raise ValueError(
        "No GeoTIFF-like asset found. Inspect the asset list and choose manually."
    )

vantor_asset_key = image_asset_keys[0]
print("Using asset:", vantor_asset_key)

vantor_subset = save_single_asset_subset(
    item=vantor_item,
    asset_key=vantor_asset_key,
    bbox_wgs84=vantor_bbox,
    output_path="data/raw/vantor/vantor_bordeaux_event_rgb_subset.tif",
    band_indexes=(1, 2, 3),
)

plot_rgb(vantor_subset, "Vantor Open Data subset: Bordeaux wildfire event")

import requests
from collections import defaultdict

# WGS84 bounding box around UZH Campus Irchel
irchel_bbox_wgs84 = [8.550, 47.3970, 8.5535, 47.3995]

# Swiss orthophotos are distributed in LV95 (EPSG:2056); transform the
# bbox once so rasterio.merge can mosaic tiles to this exact extent.
transformer = Transformer.from_crs("EPSG:4326", "EPSG:2056", always_xy=True)
minx, miny = transformer.transform(*irchel_bbox_wgs84[:2])
maxx, maxy = transformer.transform(*irchel_bbox_wgs84[2:])
irchel_bounds_lv95 = (minx, miny, maxx, maxy)

collection_id = "ch.swisstopo.swissimage-dop10"
search_url = "https://data.geo.admin.ch/api/stac/v0.9/search"

response = requests.post(
    search_url,
    json={"collections": [collection_id], "bbox": irchel_bbox_wgs84, "limit": 100},
)
response.raise_for_status()
swissimage_items = response.json().get("features", [])

# SWISSIMAGE tiles are flown in different years; group all tiles that
# intersect the bbox by year so each year can be mosaicked separately.
yearly_tiles = defaultdict(list)
for item in swissimage_items:
    year = item.get("properties", {}).get("datetime", "").split("-")[0] or "Unknown"
    for asset in item.get("assets", {}).values():
        if "image/tiff" in asset.get("type", ""):
            yearly_tiles[year].append(asset.get("href"))
            break

sorted_years = sorted(yearly_tiles)
print(f"Found imagery across {len(sorted_years)} years: {sorted_years}")
for year in sorted_years:
    print(f"  {year}: {len(yearly_tiles[year])} tile(s) to mosaic")

from rasterio.merge import merge

fig, axes = plt.subplots(1, len(sorted_years), figsize=(7 * len(sorted_years), 7))
if len(sorted_years) == 1:
    axes = [axes]

for ax, year in zip(axes, sorted_years):
    try:
        src_files = [rasterio.open(url) for url in yearly_tiles[year]]
        mosaic, _ = merge(src_files, bounds=irchel_bounds_lv95, indexes=[1, 2, 3])
        for src in src_files:
            src.close()

        rgb = np.clip(np.transpose(mosaic, (1, 2, 0)).astype(float) / 255.0, 0, 1)
        ax.imshow(rgb)
        ax.set_title(f"UZH Campus Irchel ({year})", fontweight="bold")
    except Exception as e:
        ax.text(
            0.5,
            0.5,
            f"Error: {e}",
            ha="center",
            va="center",
            transform=ax.transAxes,
            color="red",
        )
    ax.axis("off")

plt.tight_layout()
plt.suptitle("SWISSIMAGE 10cm mosaicked time series: UZH Campus Irchel", y=1.03)
plt.show()

def inspect_raster(path):
    with rasterio.open(path) as src:
        print("Path:", path)
        print("Width x height:", src.width, "x", src.height)
        print("Bands:", src.count)
        print("CRS:", src.crs)
        print("Resolution:", src.res)
        print("Bounds:", src.bounds)
        print("Data type:", src.dtypes[0])
        print("File size:", f"{Path(path).stat().st_size / 1_000_000:.2f} MB")

inspect_raster(s2_subset)

query_filter = {"eo:cloud_cover": {"lt": 20}}
bands = ["B04", "B03", "B02"]

for band in bands:
    print("Requested band:", band)

project_area = "rome"
bbox_wgs84 = [12.43, 41.88, 12.49, 41.91]
date_range = "2025-06-01/2025-08-31"
max_cloud_cover = 20