SST Gap-Filling: Does the WGS84 Ellipsoid Improve Results?

import os
import warnings
import time
import json

import numpy as np
import healpy as hp
import matplotlib
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)

DATE = "2026-04-01"

CI_MODE = os.environ.get("CI_MODE", "0") == "1"

# LEVEL is the GRID4EARTH 'depth' (nside = 2**LEVEL): nside=16 -> depth 4, nside=32 -> depth 5.
if CI_MODE:
    NSIDE = 16
    LEVEL = 4       # HEALPix level / GRID4EARTH depth for healpix-resample (2^level = NSIDE)
    NSTEPS = 20
    LMAX = 15
    print("Running in CI mode (low resolution)")
else:
    NSIDE = 32
    LEVEL = 5
    NSTEPS = 300
    LMAX = 30
    print("Running in full-resolution mode")

NPIX = hp.nside2npix(NSIDE)
print(f"NSIDE={NSIDE}, LEVEL={LEVEL}, NPIX={NPIX}, NSTEPS={NSTEPS}, LMAX={LMAX}")

import copernicusmarine

def load_sst_data(date_str):
    """Download L3S and L4 SST for a single day."""
    date_start = date_str
    date_end = date_str

    # L3S -- gappy PMW observations
    ds_l3s = copernicusmarine.open_dataset(
        dataset_id="cmems_obs-sst_glo_phy_l3s_pmw_P1D-m",
        variables=["sea_surface_temperature"],
        minimum_longitude=-180,
        maximum_longitude=180,
        minimum_latitude=-90,
        maximum_latitude=90,
        start_datetime=f"{date_start}T00:00:00",
        end_datetime=f"{date_end}T23:59:59",
    )

    # L4 -- gap-free reference
    ds_l4 = copernicusmarine.open_dataset(
        dataset_id="cmems_obs-sst_glo_phy-temp_nrt_P1D-m",
        variables=["analysed_sst"],
        minimum_longitude=-180,
        maximum_longitude=180,
        minimum_latitude=-90,
        maximum_latitude=90,
        start_datetime=f"{date_start}T00:00:00",
        end_datetime=f"{date_end}T23:59:59",
    )

    return ds_l3s, ds_l4


ds_l3s, ds_l4 = load_sst_data(DATE)
print(f"L3S shape: {ds_l3s['sea_surface_temperature'].shape}")
print(f"L4  shape: {ds_l4['analysed_sst'].shape}")

import xarray as xr

sst_l4_2d = ds_l4["analysed_sst"].isel(time=0).values  # (lat, lon)
mask_l4 = ~np.isnan(sst_l4_2d)

# Check if grids match — handle both lat/lon and latitude/longitude naming
lat_name_l3s = "lat" if "lat" in ds_l3s.dims else "latitude"
lon_name_l3s = "lon" if "lon" in ds_l3s.dims else "longitude"
lat_name_l4 = "lat" if "lat" in ds_l4.dims else "latitude"
lon_name_l4 = "lon" if "lon" in ds_l4.dims else "longitude"

l3s_lat = ds_l3s[lat_name_l3s].values
l4_lat = ds_l4[lat_name_l4].values
l3s_lon = ds_l3s[lon_name_l3s].values
l4_lon = ds_l4[lon_name_l4].values

if l3s_lat.shape != l4_lat.shape or not np.allclose(l3s_lat, l4_lat, atol=0.01):
    print("Grids differ -- regridding L3S to L4 grid via interp")
    ds_l3s_regrid = ds_l3s.interp({lat_name_l3s: l4_lat, lon_name_l3s: l4_lon}, method="nearest")
    sst_l3s_2d = ds_l3s_regrid["sea_surface_temperature"].isel(time=0).values
else:
    print("Grids match")
    sst_l3s_2d = ds_l3s["sea_surface_temperature"].isel(time=0).values

print(f"L3S 2D: {sst_l3s_2d.shape}, NaN fraction: {np.isnan(sst_l3s_2d).mean():.2%}")
print(f"L4  2D: {sst_l4_2d.shape},  NaN fraction: {np.isnan(sst_l4_2d).mean():.2%}")

from healpix_resample import GroupByResampler


def resample_to_healpix(data_2d, mask_2d, lat, lon, ellipsoid="sphere"):
    """
    Resample a 2-D lat/lon array to HEALPix.

    Parameters
    ----------
    data_2d : ndarray (nlat, nlon)
        SST values; NaN where missing.
    mask_2d : ndarray (nlat, nlon)
        Boolean ocean mask (True = ocean).
    lat, lon : 1-D arrays
        Latitude and longitude vectors.
    ellipsoid : str
        "sphere" or "WGS84".

    Returns
    -------
    hp_map : ndarray (npix,)
        HEALPix map with NaN for empty cells.
    """
    lon_grid, lat_grid = np.meshgrid(lon, lat)
    valid = mask_2d & ~np.isnan(data_2d)

    lon_deg = lon_grid[valid].ravel()
    lat_deg = lat_grid[valid].ravel()
    values = data_2d[valid].ravel()

    resampler = GroupByResampler(
        lon_deg=lon_deg,
        lat_deg=lat_deg,
        level=LEVEL,
        reduce="mean",
        ellipsoid=ellipsoid,
    )
    result = resampler.resample(values)

    # Build full-size map
    hp_map = np.full(NPIX, np.nan)
    hp_map[result.cell_ids] = result.cell_data

    return hp_map

import foscat.scat_cov as sc
import foscat.Synthesis as synth


def run_foscat(ellipsoid="sphere"):
    """
    Full gap-filling pipeline for a given ellipsoid.

    Returns
    -------
    result : dict
        Keys: ellipsoid, rmse_mk, time_s, hp_l4, hp_l3s, hp_filled, gap_mask
    """
    print(f"\n{'='*60}")
    print(f"  Running FOSCAT with ellipsoid = {ellipsoid}")
    print(f"{'='*60}")
    t0 = time.time()

    # --- 1. Resample to HEALPix ---
    lat_name = "lat" if "lat" in ds_l4.dims else "latitude"
    lon_name = "lon" if "lon" in ds_l4.dims else "longitude"
    lat = ds_l4[lat_name].values
    lon = ds_l4[lon_name].values

    hp_l4 = resample_to_healpix(sst_l4_2d, mask_l4, lat, lon, ellipsoid=ellipsoid)
    hp_l3s = resample_to_healpix(sst_l3s_2d, mask_l4, lat, lon, ellipsoid=ellipsoid)

    ocean = ~np.isnan(hp_l4)
    observed = ~np.isnan(hp_l3s)
    clouds = ocean & ~observed  # gap cells (ocean but not observed)

    n_ocean = ocean.sum()
    n_obs = observed.sum()
    n_gap = clouds.sum()
    print(f"  Ocean cells:    {n_ocean}")
    print(f"  Observed cells: {n_obs}")
    print(f"  Gap cells:      {n_gap} ({100*n_gap/n_ocean:.1f}%)")

    # --- 2. Spherical-harmonics baseline for gaps ---
    ocean_mean = np.nanmean(hp_l4[ocean])

    hp_start = hp_l3s.copy()
    # Fill NaN with ocean mean for alm computation
    hp_for_alm = np.where(np.isnan(hp_start), ocean_mean, hp_start)

    alm = hp.map2alm(hp_for_alm, lmax=LMAX)
    baseline = hp.alm2map(alm, NSIDE, verbose=False)

    # Use baseline to fill gap cells as initial guess
    hp_start[clouds] = baseline[clouds]
    # Set non-ocean cells to ocean mean (not 0 — SST is ~270-304 K)
    hp_start[~ocean] = ocean_mean
    hp_start[np.isnan(hp_start)] = ocean_mean

    print(f"  Spherical-harmonics baseline computed (LMAX={LMAX})")

    # --- 3. FOSCAT synthesis ---
    # Prepare L4 reference: fill non-ocean with ocean mean
    hp_l4_clean = hp_l4.copy()
    hp_l4_clean[~ocean | np.isnan(hp_l4_clean)] = ocean_mean

    scat_op = sc.funct(NORIENT=4, KERNELSZ=3, all_type="float32", silent=True)

    data = hp_start.reshape(1, NPIX).astype(np.float32)
    mask_ocean = ocean.reshape(1, NPIX).astype(np.float32)

    # Reference from L4 gap-free product
    ref, sref = scat_op.eval(
        hp_l4_clean.reshape(1, NPIX).astype(np.float32),
        mask=mask_ocean, calc_var=True
    )

    def The_loss(x, scat_operator, args):
        ref, mask, sref = args[0], args[1], args[2]
        learn = scat_operator.eval(x, mask=mask)
        loss = scat_operator.reduce_mean(scat_operator.square((ref - learn) / sref))
        return loss

    mask_const = scat_op.backend.constant(scat_op.backend.bk_cast(mask_ocean))
    loss = synth.Loss(The_loss, scat_op, ref, mask_const, sref)
    sy = synth.Synthesis([loss])

    mask_clouds_t = scat_op.backend.bk_cast(clouds.reshape(1, NPIX).astype(np.float32))

    omap = sy.run(scat_op.backend.bk_cast(data), EVAL_FREQUENCY=max(NSTEPS//10, 1),
                  grd_mask=mask_clouds_t, NUM_EPOCHS=NSTEPS, do_lbfgs=True)
    hp_filled = np.array(omap).ravel() if not hasattr(omap, 'numpy') else omap.numpy().ravel()
    # Mask land cells as NaN for clean visualisation
    hp_filled[~ocean] = np.nan

    elapsed = time.time() - t0

    # --- 4. RMSE on gap cells ---
    diff = hp_filled[clouds] - hp_l4[clouds]
    rmse = np.sqrt(np.nanmean(diff ** 2))
    rmse_mk = rmse * 1000  # convert K to mK

    print(f"  RMSE on gaps: {rmse_mk:.1f} mK")
    print(f"  Elapsed time: {elapsed:.1f} s")

    return {
        "ellipsoid": ellipsoid,
        "rmse_mk": float(rmse_mk),
        "time_s": float(elapsed),
        "n_ocean": int(n_ocean),
        "n_observed": int(n_obs),
        "n_gaps": int(n_gap),
        "hp_l4": hp_l4,
        "hp_l3s": hp_l3s,
        "hp_filled": hp_filled,
        "gap_mask": clouds,
    }

r_sphere = run_foscat("sphere")
r_wgs84 = run_foscat("WGS84")

print("\n" + "=" * 60)
print("  Results Summary")
print("=" * 60)
print(f"{'Geometry':<12} {'RMSE (mK)':>10} {'Time (s)':>10} {'Gaps':>8}")
print("-" * 44)
print(f"{'sphere':<12} {r_sphere['rmse_mk']:>10.1f} {r_sphere['time_s']:>10.1f} {r_sphere['n_gaps']:>8d}")
print(f"{'WGS84':<12} {r_wgs84['rmse_mk']:>10.1f} {r_wgs84['time_s']:>10.1f} {r_wgs84['n_gaps']:>8d}")
print("-" * 44)

diff_mk = r_sphere["rmse_mk"] - r_wgs84["rmse_mk"]
if diff_mk > 0:
    print(f"WGS84 improves RMSE by {diff_mk:.1f} mK ({100*diff_mk/r_sphere['rmse_mk']:.1f}%)")
elif diff_mk < 0:
    print(f"Sphere is better by {-diff_mk:.1f} mK ({100*(-diff_mk)/r_wgs84['rmse_mk']:.1f}%)")
else:
    print("No difference between sphere and WGS84")

results_dict = {
    "date": DATE,
    "nside": NSIDE,
    "level": LEVEL,
    "nsteps": NSTEPS,
    "lmax": LMAX,
    "ci_mode": CI_MODE,
    "sphere": {
        "rmse_mk": r_sphere["rmse_mk"],
        "time_s": r_sphere["time_s"],
        "n_ocean": r_sphere["n_ocean"],
        "n_observed": r_sphere["n_observed"],
        "n_gaps": r_sphere["n_gaps"],
    },
    "wgs84": {
        "rmse_mk": r_wgs84["rmse_mk"],
        "time_s": r_wgs84["time_s"],
        "n_ocean": r_wgs84["n_ocean"],
        "n_observed": r_wgs84["n_observed"],
        "n_gaps": r_wgs84["n_gaps"],
    },
    "diff_mk": float(diff_mk),
}

results_path = os.path.join("results", "comparison_results.json")
os.makedirs("results", exist_ok=True)
with open(results_path, "w") as f:
    json.dump(results_dict, f, indent=2)
print(f"Results saved to {results_path}")

# Common color range from L4 reference
vmin = np.nanmin(r_sphere["hp_l4"][r_sphere["gap_mask"]])
vmax = np.nanmax(r_sphere["hp_l4"][r_sphere["gap_mask"]])

maps = [
    (r_sphere["hp_l3s"], "L3S SST (with cloud gaps)"),
    (r_sphere["hp_filled"], f"FOSCAT Sphere — RMSE = {r_sphere['rmse_mk']:.1f} mK"),
    (r_wgs84["hp_filled"], f"FOSCAT WGS84 — RMSE = {r_wgs84['rmse_mk']:.1f} mK"),
    (r_sphere["hp_l4"], "L4 Reference (gap-free)"),
]

for data, title in maps:
    hp.mollview(data, title=title, cmap="RdYlBu_r",
                min=vmin, max=vmax, nest=True, unit="K")
    plt.savefig(os.path.join("results", f"map_{title.split()[0].lower()}.png"),
                dpi=150, bbox_inches="tight")
    plt.show()

fig, ax = plt.subplots(figsize=(6, 4))
labels = ["Sphere", "WGS84"]
rmses = [r_sphere["rmse_mk"], r_wgs84["rmse_mk"]]
colors = ["#4878CF", "#D65F5F"]

bars = ax.bar(labels, rmses, color=colors, width=0.5, edgecolor="black", linewidth=0.8)

for bar, val in zip(bars, rmses):
    ax.text(
        bar.get_x() + bar.get_width() / 2,
        bar.get_height() + 0.5,
        f"{val:.1f}",
        ha="center",
        va="bottom",
        fontsize=12,
        fontweight="bold",
    )

ax.set_ylabel("RMSE (mK)", fontsize=12)
ax.set_title(f"SST Gap-Filling RMSE: Sphere vs WGS84\n{DATE}, NSIDE={NSIDE}", fontsize=13)
ax.set_ylim(0, max(rmses) * 1.25)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)

fig.tight_layout()
chart_path = os.path.join("results", "comparison_rmse.png")
fig.savefig(chart_path, dpi=150)
plt.show()
print(f"Chart saved to {chart_path}")