"""Methane plume detection and quantification in Python with NASA EMIT imaging-spectrometer data:
three Dallas-Fort Worth landfill plumes on 25 August 2023.

Loads the EMIT methane enhancement (ppm m) from NASA's open VEDA STAC and raster API (no login),
segments the plumes against the background noise, computes the integrated mass enhancement (IME),
estimates emission rates with ERA5 10 m winds from Google's ARCO-ERA5 Zarr store, and compares
the snapshots with the facilities' 2023 EPA Greenhouse Gas Reporting Program totals.

TechnicalGIS tutorial: https://www.technicalgis.com/tutorials/methane-plume-detection-emit-python/
Data: NASA EMIT via NASA VEDA (CC0-1.0); NASA JPL EMIT plume-complex metadata (public domain);
Copernicus Sentinel-2 via Microsoft Planetary Computer; ERA5 (Copernicus C3S, CC-BY-4.0) via
Google ARCO-ERA5; US EPA GHGRP via Envirofacts (public domain).
"""
import json
import os
import sys
import warnings
from importlib.metadata import version
from pathlib import Path

import numpy as np
import pandas as pd
import requests
import rasterio
import rioxarray  # noqa: F401  (registers the .rio accessor)
import xarray as xr
import pystac_client
import planetary_computer
from pyproj import Geod, Transformer
from scipy import ndimage
from scipy.stats import norm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from shapely import contains_xy
from shapely.geometry import box, shape

warnings.filterwarnings("ignore")
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})

DATA = Path("data"); DATA.mkdir(exist_ok=True)
OUT = Path("output"); OUT.mkdir(exist_ok=True)

VEDA_STAC = "https://openveda.cloud/api/stac"
VEDA_RASTER = "https://openveda.cloud/api/raster"
JPL_PLUMES = "https://earth.jpl.nasa.gov/emit-mmgis-lb/Missions/EMIT/Layers/coverage/combined_plume_metadata.json"
PC_STAC = "https://planetarycomputer.microsoft.com/api/stac/v1"
ARCO = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
GHGRP = "https://data.epa.gov/efservice"
UTM = "EPSG:32614"          # UTM zone 14N, metres
M_CH4 = 0.01604             # kg/mol
R_GAS = 8.314               # J/(mol K)
GWP_RY2023 = 25             # GHGRP reporting year 2023 used the IPCC AR4 GWP for methane

# The three landfills in the scene, with their EPA GHGRP facility IDs and reported locations.
SITES = {
    "McCommas Bluff": {"ghgrp": 1006817, "lat": 32.67607, "lon": -96.73436},
    "Skyline (Ferris)": {"ghgrp": 1007852, "lat": 32.54992, "lon": -96.66561},
    "Camelot (Lewisville)": {"ghgrp": 1007780, "lat": 33.043321, "lon": -96.945276},
}
COLOURS = {"McCommas Bluff": "#d7301f", "Skyline (Ferris)": "#2b8cbe", "Camelot (Lewisville)": "#6a51a3"}
K_SIGMAS = (1.5, 2.0, 3.0)

geod = Geod(ellps="WGS84")
to_utm = Transformer.from_crs("EPSG:4326", UTM, always_xy=True)


def section(title):
    print(f"\n=== {title} ===")


def get(url, path, **params):
    if not path.exists():
        r = requests.get(url, params=params or None, timeout=600, headers={"User-Agent": "TechnicalGIS-tutorial"})
        r.raise_for_status()
        path.write_bytes(r.content)
    return path


section("versions")
print(f"python {sys.version.split()[0]} | rasterio {rasterio.__version__} | xarray {xr.__version__} | "
      f"numpy {np.__version__} | scipy {version('scipy')} | pystac-client {version('pystac-client')}")

# ---------------------------------------------------------------- Step 1: find EMIT methane data in the VEDA STAC
section("Step 1: EMIT methane enhancement from NASA VEDA")
veda = pystac_client.Client.open(VEDA_STAC)
col = veda.get_collection("landfill-emit")
print(f"collection: {col.id} | {col.title} | licence {col.license}")
items = sorted(veda.search(collections=[col.id]).items(), key=lambda i: i.datetime)
for it in items:
    st = it.assets["cog_default"].extra_fields["raster:bands"][0]["statistics"]
    print(f"  {it.id}: {it.datetime:%Y-%m-%d} | shape {it.assets['cog_default'].extra_fields['proj:shape']} | "
          f"bbox {[round(b, 3) for b in it.bbox]} | {st['valid_percent']:.1f}% valid")

item = next(i for i in items if i.id == "EMIT_landfill_cog_2023-08-25")
h, w = item.assets["cog_default"].extra_fields["proj:shape"]
# The COG sits in a bucket that is not public, but VEDA's TiTiler service streams it as a float32
# GeoTIFF. max_size >= the native shape returns every pixel at full resolution; band 2 is the valid mask.
tif = get(f"{VEDA_RASTER}/collections/{item.collection_id}/items/{item.id}/preview.tif",
          DATA / f"{item.id}.tif", assets="cog_default", max_size=max(h, w))
with rasterio.open(tif) as src:
    ch4 = src.read(1).astype("float64")
    valid = (src.read(2) > 0) & (ch4 > -9000)
    transform, bounds = src.transform, src.bounds
    print(f"read {tif.name}: {src.width} x {src.height} px, {src.crs.to_string()}, "
          f"pixel {src.res[0]:.6f} deg | native shape matched: {(src.height, src.width) == (h, w)}")
ch4[~valid] = np.nan
same = valid[:, 1:] & valid[:, :-1]
print(f"{np.mean(ch4[:, 1:][same] == ch4[:, :-1][same]):.0%} of horizontally adjacent valid pixels hold identical "
      "values (EMIT's orthorectification repeats ~60 m sensor pixels on the output grid)")
rows, cols = np.indices(ch4.shape)
lon = transform.c + (cols + 0.5) * transform.a
lat = transform.f + (rows + 0.5) * transform.e
ex, ny = to_utm.transform(lon, lat)

# Area of every pixel on the ellipsoid (pixels are 0.000542 deg, about 51 m x 60 m here).
dlon, dlat = transform.a, -transform.e
row_area = np.array([abs(geod.polygon_area_perimeter([0, dlon, dlon, 0],
                                                     [y - dlat / 2, y - dlat / 2, y + dlat / 2, y + dlat / 2])[0])
                     for y in lat[:, 0]])
pix_area = np.repeat(row_area[:, None], ch4.shape[1], axis=1)
print(f"valid pixels: {valid.sum():,} ({valid.mean():.1%} of the window) | pixel area {pix_area.mean():,.0f} m2 "
      f"(~{np.sqrt(pix_area.mean()):.0f} m)")
print(f"enhancement over valid pixels: median {np.nanmedian(ch4):.0f}, 1st pct {np.nanpercentile(ch4, 1):.0f}, "
      f"99th pct {np.nanpercentile(ch4, 99):.0f}, max {np.nanmax(ch4):.0f} ppm m")

# ---------------------------------------------------------------- Step 2: JPL plume complexes (public metadata)
section("Step 2: EMIT plume complexes from the JPL greenhouse gas portal")
meta = json.loads(get(JPL_PLUMES, DATA / "emit_combined_plume_metadata.json").read_text())
print(f"{len(meta['features']):,} plume-complex features in the JPL metadata file")
jpl = {}
for f in meta["features"]:
    p = f["properties"]
    if not p["UTC Time Observed"].startswith("2023-08-25"):
        continue
    pt_lon, pt_lat = p["Longitude of max concentration"], p["Latitude of max concentration"]
    if bounds.left <= pt_lon <= bounds.right and bounds.bottom <= pt_lat <= bounds.top:
        # The file holds a Point (the maximum) and a Polygon (the outline) per plume complex; keep the polygon.
        j = jpl.setdefault(p["Plume ID"], {"id": p["Plume ID"], "time": p["UTC Time Observed"], "lat": pt_lat,
                                           "lon": pt_lon, "max_ppmm": p["Max Plume Concentration (ppm m)"],
                                           "q_jpl": p["Emissions Rate Estimate (kg/hr)"], "geom": None})
        if f["geometry"]["type"] in ("Polygon", "MultiPolygon"):
            j["geom"] = shape(f["geometry"])
jpl = list(jpl.values())
for site, s in SITES.items():
    d = [geod.inv(s["lon"], s["lat"], j["lon"], j["lat"])[2] for j in jpl]
    s["jpl"] = j = jpl[int(np.argmin(d))]
    print(f"{site}: {j['id']} at {j['time']} | max {j['max_ppmm']:.0f} ppm m at ({j['lat']:.4f}, {j['lon']:.4f}), "
          f"{min(d) / 1000:.1f} km from the GHGRP location | JPL emission estimate: {j['q_jpl']}")
overpass = pd.Timestamp(jpl[0]["time"]).tz_localize(None)

# ---------------------------------------------------------------- Step 3: background noise
section("Step 3: background noise of the enhancement field")
# The VEDA raster is clipped to the JPL plume-complex outlines, so there is no plume-free area to sample.
# The outlines are generous, though, and most pixels inside them are noise: iterative 3-sigma clipping
# of all valid pixels recovers the background median and spread without the plume tail.
for j in jpl:
    inside = contains_xy(j["geom"], lon, lat)
    print(f"{j['id']}: outline {inside.sum():,} px, of which {(inside & valid).sum():,} carry data")
x = ch4[valid]
for it_ in range(10):
    med = np.median(x)
    sigma = 1.4826 * np.median(np.abs(x - med))
    keep = np.abs(ch4[valid] - med) < 3 * sigma
    if keep.sum() == x.size:
        break
    x = ch4[valid][keep]
print(f"sigma clipping converged after {it_ + 1} iterations: median {med:.0f} ppm m | robust sigma "
      f"(1.4826 x MAD) {sigma:.0f} ppm m | {np.mean(ch4[valid] > med + 3 * sigma):.2%} of pixels above +3 sigma, "
      f"{np.mean(ch4[valid] < med - 3 * sigma):.2%} below -3 sigma (Gaussian: 0.13% each)")
print(f"one-pixel detection limit at 3 sigma: {3 * sigma:,.0f} ppm m above background")

# ---------------------------------------------------------------- Step 4: plume masks
section("Step 4: segment each plume")
filled = np.where(valid, ch4 - med, 0.0)
weight = ndimage.gaussian_filter(valid.astype(float), 1.5)
smooth = np.where(valid, ndimage.gaussian_filter(filled, 1.5) / np.maximum(weight, 1e-6), np.nan)
sm_bg = smooth[valid & (weight > 0.6)]
sm_sigma = 1.4826 * np.median(np.abs(sm_bg - np.median(sm_bg)))
print(f"Gaussian smoothing (sigma 1.5 px) cuts the noise from {sigma:.0f} to {sm_sigma:.0f} ppm m")


def plume_origin(site, search=1000):
    """The smoothed-field maximum within `search` metres of JPL's maximum-concentration pixel.
    A single noisy pixel is a poor anchor; the smoothed peak is a steadier estimate of the source."""
    j = SITES[site]["jpl"]
    sx, sy = to_utm.transform(j["lon"], j["lat"])
    cand = np.where(np.hypot(ex - sx, ny - sy) <= search, np.nan_to_num(smooth, nan=-1e9), -1e9)
    r, c = np.unravel_index(np.argmax(cand), cand.shape)
    return float(lon[r, c]), float(lat[r, c]), float(smooth[r, c])


def plume_mask(k_sigma, seed_lonlat, radius=300):
    """Threshold the smoothed field at k*sigma and keep every connected component that reaches within
    `radius` metres of the plume origin. Lower thresholds therefore always give larger, nested masks."""
    hot = np.nan_to_num(smooth) > k_sigma * sm_sigma
    lab, n = ndimage.label(hot, structure=np.ones((3, 3)))
    sx, sy = to_utm.transform(*seed_lonlat)
    near = np.hypot(ex - sx, ny - sy) <= radius
    keep = np.setdiff1d(np.unique(lab[near & hot]), [0])
    mask = np.isin(lab, keep)
    # close small holes so noisy pixels inside the plume still count towards the mass
    return ndimage.binary_closing(mask, iterations=2) & valid


def kg_per_m2_per_ppmm(p_pa, t_k):
    """kg of CH4 per m2 for 1 ppm m of enhancement, for air at pressure p_pa and temperature t_k."""
    return 1e-6 * p_pa / (R_GAS * t_k) * M_CH4


# ---------------------------------------------------------------- Step 5: ERA5 wind at overpass
section("Step 5: ERA5 wind, pressure and temperature at the overpass")
era = xr.open_zarr(ARCO, chunks=None, storage_options={"token": "anon"})
vars_ = ["10m_u_component_of_wind", "10m_v_component_of_wind", "100m_u_component_of_wind",
         "100m_v_component_of_wind", "surface_pressure", "2m_temperature", "boundary_layer_height"]
day = overpass.floor("D")
sub = era[vars_].sel(time=slice(day + pd.Timedelta(hours=12), day + pd.Timedelta(hours=23)),
                     latitude=slice(33.5, 32.0), longitude=slice(262.5, 263.75)).load()
print(f"ERA5 subset: {dict(sub.sizes)} | overpass {overpass:%Y-%m-%d %H:%M} UTC "
      f"({(overpass - pd.Timedelta(hours=5)):%H:%M} local CDT)")
wind = {}
for site, s in SITES.items():
    j = s["jpl"]
    p = sub.interp(latitude=j["lat"], longitude=j["lon"] % 360)
    ts = p.interp(time=overpass)
    u10, v10 = float(ts["10m_u_component_of_wind"]), float(ts["10m_v_component_of_wind"])
    u100, v100 = float(ts["100m_u_component_of_wind"]), float(ts["100m_v_component_of_wind"])
    spd = np.hypot(p["10m_u_component_of_wind"], p["10m_v_component_of_wind"]).to_series()
    dirn = ((np.degrees(np.arctan2(-p["10m_u_component_of_wind"], -p["10m_v_component_of_wind"])) + 360) % 360)
    win = spd.loc[overpass - pd.Timedelta(hours=2): overpass + pd.Timedelta(hours=2)]
    wind[site] = wv = {"u10": u10, "v10": v10, "U10": np.hypot(u10, v10), "U100": np.hypot(u100, v100),
                       "dir_from": (np.degrees(np.arctan2(-u10, -v10)) + 360) % 360,
                       "U10_lo": float(win.min()), "U10_hi": float(win.max()),
                       "ps": float(ts["surface_pressure"]), "t2m": float(ts["2m_temperature"]),
                       "blh": float(ts["boundary_layer_height"]), "hourly": spd, "dir_hourly": dirn.to_series()}
    print(f"{site}: U10 {wv['U10']:.2f} m/s from {wv['dir_from']:.0f} deg | U100 {wv['U100']:.2f} m/s | "
          f"U10 range over +-2 h {wv['U10_lo']:.2f}-{wv['U10_hi']:.2f} m/s | ps {wv['ps'] / 100:.0f} hPa | "
          f"T2m {wv['t2m'] - 273.15:.1f} C | BLH {wv['blh']:.0f} m")

# ---------------------------------------------------------------- Step 6: IME and emission rates
section("Step 6: integrated mass enhancement and emission rate")
results, fetch_curves, masks = [], {}, {}
for site, s in SITES.items():
    j, wv = s["jpl"], wind[site]
    kfac = kg_per_m2_per_ppmm(wv["ps"], wv["t2m"])
    o_lon, o_lat, o_val = s["origin"] = plume_origin(site)
    seed = (o_lon, o_lat)
    print(f"{site}: plume origin (smoothed peak) at ({o_lat:.4f}, {o_lon:.4f}), {o_val:,.0f} ppm m above background, "
          f"{geod.inv(j['lon'], j['lat'], o_lon, o_lat)[2]:.0f} m from the JPL maximum pixel")
    sx, sy = to_utm.transform(*seed)
    dist = np.hypot(ex - sx, ny - sy)
    ueff = 1.1 * np.log(wv["U10"]) + 0.6          # Varon et al. (2018) effective wind for L = sqrt(A)
    per_k = {}
    for k in K_SIGMAS:
        m = plume_mask(k, seed)
        ime = np.nansum(np.where(m, ch4 - med, 0) * pix_area) * kfac      # kg of excess CH4 in the plume
        area = (pix_area * m).sum()
        length = dist[m].max()
        per_k[k] = {"mask": m, "ime": ime, "area": area, "length": length,
                    "q_simple": wv["U10"] * ime / length * 3600,
                    "q_varon": ueff * ime / np.sqrt(area) * 3600 if ueff > 0 else np.nan}
    base = per_k[2.0]
    masks[site] = per_k
    # Fetch curve: IME inside a growing circle around the plume origin, divided by its radius, using the
    # 1.5 sigma mask and, for contrast, every pixel in the outline (no mask at all).
    radii = np.arange(150, 5001, 150)
    ime_r = np.array([np.nansum(np.where(per_k[1.5]["mask"] & (dist <= r), ch4 - med, 0) * pix_area) * kfac
                      for r in radii])
    ime_all = np.array([np.nansum(np.where(valid & (dist <= r), ch4 - med, 0) * pix_area) * kfac
                        for r in radii])
    q_r, q_all = (wv["U10"] * x_ / radii * 3600 for x_ in (ime_r, ime_all))
    fetch_curves[site] = (radii, ime_r, q_r, ime_all, q_all)
    i_pk = int(np.argmax(q_all))
    fetch_msg = (f"   fetch curve without a mask peaks at r = {radii[i_pk] / 1000:.2f} km: IME {ime_all[i_pk]:,.0f} kg, "
                 f"Q {q_all[i_pk]:,.0f} kg/h; at r = 5 km the IME is {ime_all[-1]:,.0f} kg")
    lo = min(per_k[k]["q_simple"] for k in K_SIGMAS) * wv["U10_lo"] / wv["U10"]
    hi = max(per_k[k]["q_simple"] for k in K_SIGMAS) * wv["U10_hi"] / wv["U10"]
    ghg = requests.get(f"{GHGRP}/pub_facts_subp_ghg_emission/facility_id/{s['ghgrp']}/year/2023/JSON",
                       timeout=120).json()
    co2e = sum(r["co2e_emission"] for r in ghg if r["gas_id"] == 2)
    q_ghgrp = co2e / GWP_RY2023 * 1000 / 8760
    results.append({"site": site, "ime_kg": base["ime"], "area_km2": base["area"] / 1e6,
                    "length_km": base["length"] / 1000, "U10": wv["U10"], "Ueff": ueff,
                    "q_simple": base["q_simple"], "q_varon": base["q_varon"], "q_lo": lo, "q_hi": hi,
                    "ghgrp_t": co2e / GWP_RY2023, "q_ghgrp": q_ghgrp})
    print(f"{site}: {kfac * 1e6:.1f} mg CH4/m2 per ppm m | 2 sigma mask {base['mask'].sum():,} px = "
          f"{base['area'] / 1e6:.2f} km2, reaching {base['length'] / 1000:.1f} km from the origin")
    print("   IME " + ", ".join(f"{k} sigma: {per_k[k]['ime']:,.0f} kg" for k in K_SIGMAS))
    print("   Q = U10*IME/L " + ", ".join(f"{k} sigma: {per_k[k]['q_simple']:,.0f}" for k in K_SIGMAS) + " kg/h")
    varon = (f"{base['q_varon']:,.0f} kg/h" if ueff > 0 else "undefined (U_eff <= 0 below U10 = 0.58 m/s)")
    print(f"   Varon U_eff {ueff:.2f} m/s -> U_eff*IME/sqrt(A) {varon} | "
          f"range with wind +-2 h and thresholds: {lo:,.0f}-{hi:,.0f} kg/h")
    print(f"   EPA GHGRP 2023: {co2e:,.0f} t CO2e of CH4 = {co2e / GWP_RY2023:,.0f} t CH4 = {q_ghgrp:,.0f} kg/h "
          f"on average")
    print(fetch_msg)
res = pd.DataFrame(results).set_index("site")
res.round(3).to_csv(OUT / "emission_estimates.csv")
print("\n" + res[["ime_kg", "length_km", "U10", "q_simple", "q_varon", "q_lo", "q_hi", "q_ghgrp"]]
      .round(1).to_string())
tot = res[["q_simple", "q_lo", "q_hi", "q_ghgrp"]].sum()
print(f"three landfills together: {tot.q_simple:,.0f} kg/h (range {tot.q_lo:,.0f}-{tot.q_hi:,.0f}) vs GHGRP "
      f"annual average {tot.q_ghgrp:,.0f} kg/h | ratio {tot.q_simple / tot.q_ghgrp:.2f}")

# ---------------------------------------------------------------- Step 7: Sentinel-2 context
section("Step 7: Sentinel-2 true colour for context (Planetary Computer)")
pc = pystac_client.Client.open(PC_STAC, modifier=planetary_computer.sign_inplace)
s2 = sorted(pc.search(collections=["sentinel-2-l2a"], bbox=list(bounds),
                      datetime="2023-08-15/2023-09-05", query={"eo:cloud_cover": {"lt": 5}}).items(),
            key=lambda i: abs((i.datetime.replace(tzinfo=None) - overpass).total_seconds()))
print(f"{len(s2)} Sentinel-2 L2A scenes with <5% cloud within about 10 days")


def s2_rgb(bbox_ll):
    """True-colour Sentinel-2 crop in EPSG:4326 for bbox (lon0, lat0, lon1, lat1)."""
    for it in s2:
        if not shape(it.geometry).contains(box(*bbox_ll)):
            continue
        bands = []
        for b in ("B04", "B03", "B02"):
            da = rioxarray.open_rasterio(it.assets[b].href).squeeze()
            bands.append(da.rio.clip_box(*bbox_ll, crs="EPSG:4326").rio.reproject("EPSG:4326"))
        rgb = np.clip(np.dstack([(b.values.astype(float) - 1000) / 2200 for b in bands]), 0, 1) ** 0.8
        return it, rgb, bands[0].rio.bounds()
    return None, None, None


# ---------------------------------------------------------------- Figures
section("Figures")


def site_window(site, pad=0.01):
    s = SITES[site]["jpl"]
    rr, cc = np.where(valid & (np.abs(lat - s["lat"]) < 0.06) & (np.abs(lon - s["lon"]) < 0.13))
    r0, r1, c0, c1 = rr.min(), rr.max() + 1, cc.min(), cc.max() + 1
    ext = [transform.c + c0 * transform.a, transform.c + c1 * transform.a,
           transform.f + r1 * transform.e, transform.f + r0 * transform.e]
    return (r0, r1, c0, c1), ext, (ext[0] - pad, ext[2] - pad, ext[1] + pad, ext[3] + pad)


fig, axes = plt.subplots(1, 3, figsize=(13, 4.6), gridspec_kw={"width_ratios": [1, 1.6, 1.6]}, layout="constrained")
s2_dates = set()
for ax, site in zip(axes, ["Camelot (Lewisville)", "McCommas Bluff", "Skyline (Ferris)"]):
    (r0, r1, c0, c1), ext, bb = site_window(site)
    it, rgb, b = s2_rgb(bb)
    if rgb is not None:
        ax.imshow(rgb, extent=[b[0], b[2], b[1], b[3]])
        s2_dates.add(f"{it.datetime:%d %b %Y}")
    enh = ch4[r0:r1, c0:c1] - med
    im = ax.imshow(np.where(masks[site][1.5]["mask"][r0:r1, c0:c1], enh, np.nan), extent=ext, cmap="plasma",
                   vmin=0, vmax=1500, interpolation="nearest")
    s = SITES[site]
    ax.plot(s["lon"], s["lat"], marker="^", ms=8, mfc="white", mec="black", ls="none")
    ax.set_xlim(bb[0], bb[2]); ax.set_ylim(bb[1], bb[3])
    ax.set_aspect(1 / np.cos(np.radians(s["lat"])))
    ax.set_title(f"{site}\nplume complex {s['jpl']['id'].split('-')[-1]}, JPL max {s['jpl']['max_ppmm']:,.0f} ppm m", fontsize=9)
    ax.tick_params(labelsize=7)
    ax.set_xlabel("Longitude")
axes[0].set_ylabel("Latitude")
cb = fig.colorbar(im, ax=axes, shrink=0.85, pad=0.01)
cb.set_label("CH$_4$ enhancement above background (ppm m)")
fig.suptitle(f"EMIT methane plumes over three Dallas-Fort Worth landfills, {overpass:%d %B %Y} {overpass:%H:%M} UTC\n"
             f"(1.5σ plume masks over Sentinel-2 true colour from {', '.join(sorted(s2_dates))}; "
             "triangles: EPA GHGRP facility locations)", fontsize=10)
fig.savefig(OUT / "01-emit-plumes-landfills.png", facecolor="white")
plt.close(fig)
print(f"Sentinel-2 context scenes used: {sorted(s2_dates)}")

fig, ax = plt.subplots(1, 2, figsize=(11, 4))
bins = np.arange(-2500, 3250, 50)
inside = np.zeros_like(valid)
for site in SITES:
    inside |= masks[site][2.0]["mask"]
ax[0].hist(ch4[valid & ~inside], bins=bins, color="#9e9ac8", label="outside the 2σ plume masks")
ax[0].hist(ch4[inside], bins=bins, color="#d7301f", alpha=0.85, label="inside the 2σ plume masks")
xg = np.linspace(-2500, 3200, 400)
ax[0].plot(xg, (valid & ~inside).sum() * 50 * norm.pdf(xg, med, sigma), color="black", lw=1.2,
           label=f"Gaussian: median {med:.0f}, σ {sigma:.0f}")
ax[0].set_xlabel("EMIT CH$_4$ enhancement (ppm m)"); ax[0].set_ylabel("pixels (log scale)")
ax[0].set_yscale("log"); ax[0].set_ylim(0.8, 5000)
ax[0].legend(fontsize=8, frameon=False)
ax[0].set_title(f"Single pixels are noisy: σ ≈ {sigma:.0f} ppm m")
qs = np.linspace(0.001, 0.999, 300)
for name, sel, c in (("outside the 2σ masks", valid & ~inside, "#2b8cbe"), ("inside the 2σ masks", inside, "#d7301f")):
    ax[1].plot(norm.ppf(qs), np.quantile(ch4[sel], qs), color=c, lw=2, label=name)
ax[1].plot(norm.ppf(qs), med + sigma * norm.ppf(qs), color="black", ls="--", lw=1, label="Gaussian noise")
ax[1].set_xlabel("standard normal quantile"); ax[1].set_ylabel("enhancement quantile (ppm m)")
ax[1].legend(fontsize=8, frameon=False)
ax[1].set_title("Q-Q plot: plume pixels sit in the upper tail")
fig.savefig(OUT / "02-noise-histogram.png", facecolor="white")
plt.close(fig)

site = "McCommas Bluff"
s, wv = SITES[site], wind[site]
(r0, r1, c0, c1), ext, _ = site_window(site, 0)
fig, ax = plt.subplots(3, 1, figsize=(8, 9.6), layout="constrained")
im0 = ax[0].imshow(ch4[r0:r1, c0:c1] - med, extent=ext, cmap="RdBu_r", vmin=-1500, vmax=1500, interpolation="nearest")
ax[0].set_title("Raw enhancement minus background median")
im = ax[1].imshow(smooth[r0:r1, c0:c1], extent=ext, cmap="RdBu_r", vmin=-500, vmax=500, interpolation="nearest")
ax[1].set_title(f"Gaussian-smoothed (σ 1.5 px): noise {sm_sigma:.0f} ppm m")
im2 = ax[2].imshow(smooth[r0:r1, c0:c1], extent=ext, cmap="Greys_r", vmin=-500, vmax=500, interpolation="nearest")
mask_cols = ("#fdae61", "#d7301f", "#7f0000")
for k, colr in zip(K_SIGMAS, mask_cols):
    ax[2].contour(masks[site][k]["mask"][r0:r1, c0:c1].astype(float), levels=[0.5], colors=[colr], linewidths=1.3,
                  extent=ext, origin="upper")
gx, gy = s["jpl"]["geom"].exterior.xy
ax[2].plot(gx, gy, color="#41ab5d", lw=1.2, ls="--")
handles = [Line2D([], [], color=c, lw=1.5, label=f"{k}σ mask") for k, c in zip(K_SIGMAS, mask_cols)]
handles.append(Line2D([], [], color="#41ab5d", lw=1.5, ls="--", label="JPL plume-complex outline"))
ax[2].legend(handles=handles, fontsize=7, loc="lower left", framealpha=0.9)
ax[2].set_title("Plume masks and the JPL outline")
ax0x, ax0y = ext[0] + 0.012, ext[3] - 0.012
for a in ax:
    a.plot(s["origin"][0], s["origin"][1], marker="*", ms=11, mfc="yellow", mec="black", ls="none")
    a.set_aspect(1 / np.cos(np.radians(s["lat"])))
    a.tick_params(labelsize=7)
    a.annotate("", xy=(ax0x + 0.012 * wv["u10"] / wv["U10"], ax0y + 0.012 * wv["v10"] / wv["U10"]),
               xytext=(ax0x, ax0y), arrowprops=dict(arrowstyle="->", color="black", lw=1.6))
    a.text(ext[0] + 0.03, ext[3] - 0.004, f"ERA5 10 m wind {wv['U10']:.1f} m/s", fontsize=7, va="top",
           bbox=dict(fc="white", ec="none", alpha=0.85))
fig.colorbar(im0, ax=ax[0], shrink=0.9, label="ppm m")
fig.colorbar(im, ax=ax[1], shrink=0.9, label="ppm m")
fig.colorbar(im2, ax=ax[2], shrink=0.9, label="ppm m (smoothed)")
fig.suptitle(f"Segmenting the McCommas Bluff plume, {overpass:%d %B %Y} (star: plume origin)", fontsize=10)
fig.savefig(OUT / "03-plume-mask-mccommas.png", facecolor="white")
plt.close(fig)

fig, ax = plt.subplots(1, 2, figsize=(11, 3.8))
for site in SITES:
    wv = wind[site]
    hrs = wv["hourly"].index.hour
    ax[0].plot(hrs, wv["hourly"].values, marker="o", ms=3, color=COLOURS[site], label=site)
    ax[1].plot(hrs, wv["dir_hourly"].values, marker="o", ms=3, color=COLOURS[site], label=site)
for a in ax:
    a.axvline(overpass.hour + overpass.minute / 60, color="grey", ls="--", lw=1)
    a.set_xlabel(f"hour (UTC) on {overpass:%d %B %Y}")
ax[0].set_ylabel("ERA5 10 m wind speed (m/s)"); ax[0].set_ylim(0, None)
ax[1].set_ylabel("direction the wind blows from (°)"); ax[1].set_ylim(0, 360); ax[1].set_yticks([0, 90, 180, 270, 360])
ax[0].legend(fontsize=8, frameon=False)
ax[0].text(overpass.hour + 0.2, ax[0].get_ylim()[1] * 0.95, "EMIT overpass", fontsize=7, va="top")
fig.suptitle("ERA5 10 m winds at the three plume origins", fontsize=10)
fig.savefig(OUT / "04-era5-wind.png", facecolor="white")
plt.close(fig)

fig, ax = plt.subplots(1, 2, figsize=(11, 3.9), layout="constrained")
for site in SITES:
    radii, ime_r, q_r, ime_all, q_all = fetch_curves[site]
    ax[0].plot(radii / 1000, ime_r, color=COLOURS[site], label=site)
    ax[0].plot(radii / 1000, ime_all, color=COLOURS[site], ls="--", lw=1)
    ax[1].plot(radii / 1000, q_r, color=COLOURS[site], label=site)
    ax[1].plot(radii / 1000, q_all, color=COLOURS[site], ls="--", lw=1)
ax[0].set_xlabel("radius from plume origin (km)"); ax[0].set_ylabel("IME inside radius (kg CH$_4$)")
ax[1].set_xlabel("fetch radius r (km)"); ax[1].set_ylabel("Q(r) = U$_{10}$ · IME(r) / r  (kg/h)")
ax[0].legend(handles=ax[0].get_legend_handles_labels()[0] + [
    Line2D([], [], color="grey", label="inside the 1.5σ mask"),
    Line2D([], [], color="grey", ls="--", lw=1, label="every pixel in the outline")], fontsize=8, frameon=False)
ax[0].set_title("Integrated mass enhancement vs distance")
ax[1].set_title("Emission estimate vs fetch")
fig.savefig(OUT / "05-ime-fetch-curve.png", facecolor="white")
plt.close(fig)

fig, ax = plt.subplots(figsize=(8, 4))
x = np.arange(len(res))
ax.bar(x - 0.2, res.q_simple, 0.38, color=[COLOURS[s] for s in res.index], label="EMIT snapshot, U$_{10}$·IME/L (2σ)")
ax.errorbar(x - 0.2, res.q_simple, yerr=[res.q_simple - res.q_lo, res.q_hi - res.q_simple], fmt="none",
            ecolor="black", capsize=4, lw=1)
ax.scatter(x - 0.2, res.q_varon, marker="D", color="black", zorder=3, s=18, label="Varon et al. U$_{eff}$·IME/√A")
ax.bar(x + 0.2, res.q_ghgrp, 0.38, color="#bdbdbd", edgecolor="#636363", hatch="//",
       label="EPA GHGRP 2023 (annual average)")
ax.set_xticks(x); ax.set_xticklabels(res.index)
ax.set_ylabel("methane emission rate (kg/h)")
ax.legend(fontsize=8, frameon=False)
ax.set_title("One EMIT snapshot vs the reported inventory (error bars: wind ±2 h and 1.5–3σ masks)", fontsize=10)
fig.savefig(OUT / "06-emission-comparison.png", facecolor="white")
plt.close(fig)
print("wrote", ", ".join(sorted(p.name for p in OUT.glob("*.png"))))
sys.stdout.flush()
os._exit(0)   # gcsfs can leave a background thread running; exit cleanly