Mapping xOPR frames to ICESat-2 ATL06 granules

import time

import fsspec
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from xopr import crossovers

# Post-#78 catalogs with opr:mbox (current as of 2026-04). Antarctic, Bedmap-prioritized.
CATALOG_URL = (
    "https://data.source.coop/englacial/xopr/catalog/"
    "hemisphere=south/provider=cresis/collection=2016_Antarctica_DC8/stac.parquet"
)

# Pick exactly one of: DATE_RANGE, CYCLE, or TEMPORAL_MODE. The other two stay None.
DATE_RANGE = None                     # e.g. ("2024-01-01", "2024-06-30")
CYCLE = 22                            # one ICESat-2 repeat cycle (91 days); magg-compatible
TEMPORAL_MODE = None                  # "exact_year" or "all_years" — used only if the above are None

CRS = "EPSG:3031"                     # Antarctic polar stereographic

# Tunable for the prefix backend — samples per boundary edge of each granule polygon.
# ATL06 granules are thin strips so 50 is plenty.
PREFIX_N_PER_SEGMENT = 50

with fsspec.open(CATALOG_URL) as f:
    frames = gpd.read_parquet(f)

if "opr:mbox" not in frames.columns:
    raise RuntimeError(
        "Catalog has no opr:mbox column — it predates PR #77 / issue #78. "
        "Pick a reprocessed season from https://github.com/englacial/xopr/issues/78."
    )

print(f"Loaded {len(frames)} frames from {frames['collection'].iloc[0]}")
print(f"Date range: {frames['opr:date'].min()} – {frames['opr:date'].max()}")
frames.head(3)

# Example: a few Antarctic points (lat, lon). Swap in your own, or set None to skip.
POINTS_OF_INTEREST = [
    (-80.5, -78.0),   # Thwaites
    (-75.1, -105.0),  # Pine Island area
    (-70.5,  70.0),   # Amery
]

if POINTS_OF_INTEREST is not None:
    n_before = len(frames)
    frames = crossovers.subset_frames_by_points(frames, POINTS_OF_INTEREST)
    print(f"Subset by {len(POINTS_OF_INTEREST)} points: {n_before} → {len(frames)} frames")
else:
    print(f"No subset applied; working with all {len(frames)} frames")

if len(frames) == 0:
    raise RuntimeError("Subset left zero frames — widen POINTS_OF_INTEREST or set it to None.")

# Flatten every frame's 4 mbox cells into one list for the bbox computation.
all_cells = [int(c) for mbox in frames["opr:mbox"] for c in mbox]
bbox = crossovers.cmr_bbox_from_mpolygon(all_cells, step=16)
print(f"CMR bbox (lon_min, lat_min, lon_max, lat_max): {bbox}")

start_date, end_date = crossovers.resolve_temporal_window(
    frames_gdf=frames,
    mode=TEMPORAL_MODE or "exact_year",
    cycle=CYCLE,
    date_range=DATE_RANGE,
)
print(f"Temporal window: {start_date} → {end_date}")

t0 = time.perf_counter()
granules = crossovers.query_cmr_granules(
    start_date=start_date,
    end_date=end_date,
    short_name="ATL06",
    version="007",
    provider="NSIDC_CPRD",
    bbox=bbox,
)
t_cmr = time.perf_counter() - t0
print(f"CMR returned {len(granules)} granules in {t_cmr:.1f}s")

t0 = time.perf_counter()
tree, st_records = crossovers.build_granule_strtree(granules, crs=CRS)
t_build_st = time.perf_counter() - t0

t0 = time.perf_counter()
strtree_result = crossovers.match_frames_to_granules(
    frames, tree, st_records, crs=CRS, step=32,
)
t_match_st = time.perf_counter() - t0

print(f"Built STRtree over {len(st_records)} granule polygons in {t_build_st:.2f}s")
print(f"Matched {len(frames)} frames in {t_match_st:.2f}s")
strtree_result[["id", "opr:date", "opr:frame", "n_granules"]].head()

strtree_result

strtree_result['atl06_granules'].iloc[0]

t0 = time.perf_counter()
pf_records = crossovers.build_granule_prefix_index(
    granules, n_per_segment=PREFIX_N_PER_SEGMENT,
)
t_build_pf = time.perf_counter() - t0

t0 = time.perf_counter()
prefix_result = crossovers.match_frames_to_granules_prefix(frames, pf_records)
t_match_pf = time.perf_counter() - t0

print(f"Built prefix index over {len(pf_records)} granules in {t_build_pf:.2f}s")
print(f"Matched {len(frames)} frames in {t_match_pf:.2f}s")
prefix_result[["id", "opr:date", "opr:frame", "n_granules"]].head()

timing = pd.DataFrame(
    {
        "stage": ["index build", "frame matching", "total local"],
        "STRtree (s)": [t_build_st, t_match_st, t_build_st + t_match_st],
        "prefix (s)": [t_build_pf, t_match_pf, t_build_pf + t_match_pf],
    }
)
timing["prefix / STRtree"] = timing["prefix (s)"] / timing["STRtree (s)"]
timing.style.format("{:.3f}", subset=["STRtree (s)", "prefix (s)", "prefix / STRtree"])

def _ids(row):
    return {g["granule_id"] for g in row}

st_sets = strtree_result["atl06_granules"].apply(_ids)
pf_sets = prefix_result["atl06_granules"].apply(_ids)

missed_by_prefix = [len(st - pf) for st, pf in zip(st_sets, pf_sets)]
extra_from_prefix = [len(pf - st) for st, pf in zip(st_sets, pf_sets)]

summary = pd.DataFrame(
    {
        "STRtree n_granules": strtree_result["n_granules"],
        "prefix n_granules": prefix_result["n_granules"],
        "missed by prefix": missed_by_prefix,
        "extra from prefix": extra_from_prefix,
    }
)
print(f"Frames with any disagreement: {(summary['missed by prefix'] + summary['extra from prefix'] > 0).sum()}/{len(summary)}")
print(f"Mean missed by prefix: {np.mean(missed_by_prefix):.2f}")
print(f"Mean extra from prefix: {np.mean(extra_from_prefix):.2f}")
summary.describe()

result = strtree_result
print(f"{len(result)} frames, {result['n_granules'].sum()} total frame↔granule edges")
print(f"Frames with ≥1 granule: {(result['n_granules'] > 0).sum()} ({100*(result['n_granules'] > 0).mean():.1f}%)")
print(f"Median granules per frame: {int(result['n_granules'].median())}")
result[["id", "opr:segment", "opr:date", "opr:frame", "geometry", "atl06_granules", "n_granules"]].head()

# Show the granules for one frame with hits
hit = result[result["n_granules"] > 0].iloc[0]
print(f"Frame {hit['id']} matches {hit['n_granules']} ATL06 granules.")
print("First three:")
for g in hit["atl06_granules"][:3]:
    print(f"  {g['granule_id']}")
    print(f"    s3:    {g['s3_url']}")
    print(f"    https: {g['https_url']}")

result_3031 = result.to_crs(CRS)
fig, ax = plt.subplots(1, 1, figsize=(9, 9))
result_3031.plot(
    column="n_granules",
    cmap="viridis",
    linewidth=0.6,
    legend=True,
    ax=ax,
)
ax.set_title(f"{frames['collection'].iloc[0]} — ATL06 hits per frame")
ax.set_xlabel("x (m, EPSG:3031)")
ax.set_ylabel("y (m, EPSG:3031)")
ax.set_aspect("equal")
plt.tight_layout()