Where the classifier flags internal waves in the SWOT archive

def scan(out: Path) -> None:
    """World map of archive footprints flagged IW, and the archive score distribution."""
    rows = read_scan(Path("runs/scan_full"))
    lon = (np.array([float(r["lon"]) for r in rows]) + 180) % 360 - 180
    lat = np.array([float(r["lat"]) for r in rows])
    score = np.array([float(r["score"]) for r in rows])
    iw = np.array([int(r["iw"]) for r in rows])
    edges = np.arange(-180, 181, 1.0), np.arange(-90, 91, 1.0)
    count, _, _ = np.histogram2d(lon, lat, edges)
    hits, _, _ = np.histogram2d(lon, lat, edges, weights=iw)
    cx, cy = np.nonzero(count >= 10)
    frac = hits[cx, cy] / count[cx, cy]
    order = np.argsort(frac)

    fig = plt.figure(figsize=(TEXTWIDTH, 2.25))
    gs = fig.add_gridspec(1, 2, width_ratios=[2.9, 1], wspace=0.4, left=0.01, right=0.99,
                          bottom=0.17, top=0.88)
    ax = fig.add_subplot(gs[0], projection=ccrs.Robinson())
    ax.set_global()
    ax.add_feature(cfeature.LAND, facecolor="#e6e6e6", edgecolor="none", zorder=0)
    ax.add_feature(cfeature.COASTLINE, edgecolor=MUTED, linewidth=0.3, zorder=1)
    ax.spines["geo"].set_edgecolor(MUTED)
    ax.spines["geo"].set_linewidth(0.5)
    dots = ax.scatter(edges[0][cx][order] + 0.5, edges[1][cy][order] + 0.5, c=frac[order], s=2.5, lw=0,
                      cmap="YlOrRd", vmin=0, vmax=0.6, transform=ccrs.PlateCarree(), rasterized=True, zorder=2)
    cbar = fig.colorbar(dots, cax=ax.inset_axes([1.01, 0.1, 0.02, 0.8]), extend="max")
    cbar.set_label("fraction flagged", labelpad=2)
    cbar.outline.set_linewidth(0.4)
    ax.set_title("(a) archive footprints flagged as internal wave", loc="left", fontsize=8.5)

    hax = fig.add_subplot(gs[1])
    style_grid(hax)
    hax.hist(score, bins=np.linspace(0, 1, 41), color="#2a78d6", lw=0)
    hax.set_yscale("log")
    hax.axvline(THRESHOLD, color=INK, lw=0.8, ls=(0, (3, 1.5)))
    hax.text(THRESHOLD + 0.02, 0.95, f"threshold {THRESHOLD}", transform=hax.get_xaxis_transform(),
             fontsize=6.5, va="top")
    hax.set_xlabel("classifier score")
    hax.set_ylabel("footprints")
    hax.set_title("(b) archive score distribution", loc="left", fontsize=8.5)
    fig.savefig(out, dpi=600)
    plt.close(fig)

"""Fit a SWOT repeat-orbit ground-track model to labeled footprints; predict lat/lon per (pass, line, side)."""
from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import numpy as np
from scipy.optimize import least_squares

MASTER = Path("/mnt/lustre/koa/koastore/ore_waves_group/data/swot/labeled/iw_master.csv")
EARTH_KM = 6371.0
E2 = 0.00669438                       # WGS84 eccentricity squared
NAMES = ["lam0_rad", "phase_asc_rad", "phase_desc_rad", "rad_per_line", "crosstrack_km",
         "incl_rad", "earth_turns_per_rev"]
# node lon, phases, 250 m lines, swath centre offset, 77.6 deg, 21-day / 292-rev repeat
START = np.array([np.radians(22.0), 0.008, -0.002, 3.81e-5, -40.0, np.radians(77.6), 21 / 292])
SWATH_CENTRE_KM = 35.0               # KaRIn swaths span 10-60 km either side of nadir
TRUSTED_KM = 10.0                     # per-pass correction kept only if its held-out median is below this
SCALE = np.array([0.01, 0.001, 0.001, 1e-7, 1.0, 0.001, 1e-4])


def load_labeled(path: Path = MASTER):
    """(pass, line index, lat, lon) of every labeled footprint; all are side R."""
    rows = []
    with open(path, encoding="utf-8") as fh:
        for row in csv.DictReader(fh):
            _, pas, idx = map(int, row["base_image_id"].split("_"))
            rows.append((pas, idx, float(row["lat"]), float(row["lon"])))
    return np.array(rows, dtype=float)


def _nadir(params, pas, idx):
    """Unit ECEF vectors of the nadir point; u counts from the south turning point of pass 1."""
    lam0, c_asc, c_desc, k, _, inc, rep = params
    u = -np.pi / 2 + np.pi * (pas - 1) + np.where(pas % 2 == 1, c_asc, c_desc) + k * idx
    x, y, z = np.cos(u), np.cos(inc) * np.sin(u), np.sin(inc) * np.sin(u)
    lon = lam0 + np.arctan2(y, x) - rep * u
    lat = np.arcsin(z)
    return np.stack([np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], -1)


def predict(params, pas, idx, side):
    """Geodetic (lat, lon) in degrees; side is +1 for R, -1 for L (swath mirrored about nadir)."""
    pas, idx, side = (np.asarray(a, dtype=float) for a in (pas, idx, side))
    here = _nadir(params, pas, idx)
    ahead = _nadir(params, pas, idx + 10.0)
    right = np.cross(ahead - here, here)
    right /= np.linalg.norm(right, axis=-1, keepdims=True)
    delta = side[..., None] * params[4] / EARTH_KM
    point = np.cos(delta) * here + np.sin(delta) * right
    lat_c = np.arcsin(np.clip(point[..., 2], -1, 1))
    lat = np.degrees(np.arctan(np.tan(lat_c) / (1 - E2)))
    lon = (np.degrees(np.arctan2(point[..., 1], point[..., 0])) + 180) % 360 - 180
    return lat, lon


def _resid(params, data):
    lat, lon = predict(params, data[:, 0], data[:, 1], np.ones(len(data)))
    dlon = (lon - data[:, 3] + 180) % 360 - 180
    return np.concatenate([(lat - data[:, 2]) * 111.2,
                           dlon * 111.2 * np.cos(np.radians(data[:, 2]))])


def error_km(params, data):
    """Per-point model-minus-label distance (local flat-earth km)."""
    return np.hypot(*_resid(params, data).reshape(2, -1))


def fit(data):
    """Robust least-squares fit of all seven orbit parameters."""
    best = None
    for lam0 in np.radians(np.arange(12.0, 32.0, 2.0)):
        res = least_squares(_resid, np.r_[lam0, START[1:]], args=(data,), loss="soft_l1",
                            f_scale=5.0, x_scale=SCALE)
        if best is None or res.cost < best.cost:
            best = res
    return best.x


def pass_phase(params, data):
    """Along-track phase correction for one pass's labeled rows."""
    def resid(x):
        return _resid(np.r_[params[0], params[1:3] + x[0], params[3:]], data)
    return float(least_squares(resid, [0.0], loss="soft_l1", f_scale=5.0).x[0])


def pass_corrections(params, data, seed=0):
    """{pass: phase} for passes whose correction holds up on a 2-fold within-pass split."""
    rng = np.random.default_rng(seed)
    out = {}
    for pas in np.unique(data[:, 0]):
        rows = rng.permutation(np.flatnonzero(data[:, 0] == pas))
        halves = rows[:len(rows) // 2], rows[len(rows) // 2:]
        errs = []
        for train, test in (halves, halves[::-1]):
            dc = pass_phase(params, data[train])
            errs.append(error_km(np.r_[params[0], params[1:3] + dc, params[3:]], data[test]))
        held = float(np.median(np.concatenate(errs)))
        print(f"  pass {int(pas):3d} corrected, held-out median {held:.2f} km")
        if held < TRUSTED_KM:
            out[int(pas)] = pass_phase(params, data[data[:, 0] == pas])
    return out


def load_fit(path: Path):
    """(parameter vector, {pass: phase correction}) from the json written by main."""
    saved = json.loads(path.read_text(encoding="utf-8"))
    return np.array([saved[n] for n in NAMES]), {int(k): v for k, v in saved["pass_phase"].items()}


def footprint_latlon(model, pas, idx, side):
    """Scene-centre and swath-centre (lat, lon) for arrays of pass, line index and side ('L'/'R')."""
    params, phase = model
    pas, idx = np.asarray(pas, dtype=float), np.asarray(idx, dtype=float)
    dc = np.array([phase.get(int(p), 0.0) for p in pas])
    shifted = idx + dc / params[3]
    centre = predict(params, pas, shifted, np.ones(len(pas)))
    sign = np.where(np.asarray(side) == "R", 1.0, -1.0)
    swath = np.r_[params[:4], [SWATH_CENTRE_KM], params[5:]]
    return centre, predict(swath, pas, shifted, sign)


def main() -> None:
    """Fit on all labeled rows, report leave-one-pass-out error, write the parameters."""
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out", type=Path, required=True, help="parameters json")
    opts = ap.parse_args()
    data = load_labeled()
    params = fit(data)
    err = error_km(params, data)
    print(f"in-sample km: median {np.median(err):.2f}  p95 {np.percentile(err, 95):.2f}  max {err.max():.0f}")
    held = []
    for pas in np.unique(data[:, 0]):
        mask = data[:, 0] == pas
        held.append(error_km(fit(data[~mask]), data[mask]))
        print(f"  pass {int(pas):3d} held out: n {mask.sum():3d}  median {np.median(held[-1]):.2f} km")
    held = np.concatenate(held)
    print(f"leave-pass-out km: median {np.median(held):.2f}  p95 {np.percentile(held, 95):.2f}  "
          f"max {held.max():.0f}")
    phase = pass_corrections(params, data)
    print(f"per-pass corrections kept for {sorted(phase)}")
    saved = dict(zip(NAMES, map(float, params))) | {"pass_phase": phase}
    opts.out.write_text(json.dumps(saved, indent=1), encoding="utf-8")
    print(dict(zip(NAMES, np.round(params, 8))))


if __name__ == "__main__":
    main()

"""Run the frozen IW classifier over every footprint of one pass in full_data_v1.h5 and write a CSV with lat/lon."""
from __future__ import annotations

import argparse
import csv
import logging
import time
from pathlib import Path

import h5py
import numpy as np
from PIL import Image

from src.data.build_channel_variant import transform
from src.data.footprint_latlon import MASTER, footprint_latlon, load_fit
from src.viz.visualize_boxes import load_model
from src.viz.visualize_compare import predict_boxes_letterboxed_with_class

LOGGER = logging.getLogger(__name__)
H5 = Path("/mnt/lustre/koa/koastore/ore_waves_group/data/swot/full_data_v1.h5")
FINAL = Path("runs/rfdetr/final/cls_iwfp_dropssh_fusion3")
CKPTS = ("s1_checkpoint_89.ckpt", "s2_checkpoint_39.ckpt", "s3_checkpoint_79.ckpt")
THRESHOLD = 0.325
MIN_SCORE = 0.01
FIELDS = ["group", "cycle", "pass", "line", "side", "valid_frac", "score_s1", "score_s2", "score_s3",
          "score", "iw", "box_x0", "box_y0", "box_x1", "box_y1", "lat", "lon", "swath_lat", "swath_lon",
          "labeled"]


def footprint_image(stacked) -> tuple[Image.Image, float]:
    """The classifier's dropssh RGB input from a (3, 660, 220) stack, and the non-NaN fraction of raw."""
    valid = float(np.isfinite(stacked[0]).mean())
    arr = np.clip(np.transpose(stacked, (1, 2, 0)), 0, 255).astype("uint8")
    return Image.fromarray(transform(arr, "dropssh")), valid


def top_box(models, image):
    """Per-model max IW score and the box of the single highest-scoring detection."""
    scores, best = [], (0.0, None)
    for model in models:
        preds = predict_boxes_letterboxed_with_class(model, image, MIN_SCORE)
        top = max(preds, key=lambda p: p[1], default=None)
        scores.append(float(top[1]) if top else 0.0)
        if top and top[1] > best[0]:
            best = (float(top[1]), top[0])
    return scores, best[1]


def labeled_groups(path: Path = MASTER) -> set[str]:
    """h5 groups (Pass_xxx/scene/side) that carry human labels."""
    with open(path, encoding="utf-8") as fh:
        return {row["h5_group"] for row in csv.DictReader(fh)}


def footprint_row(h5, group: str, models, fit, labeled: set[str]) -> dict:
    """Classify one footprint and locate it."""
    # pylint: disable=too-many-locals
    _, scene, side = group.split("/")
    cycle, pas, line = map(int, scene.split("_"))
    image, valid = footprint_image(h5[group]["stacked"][...])
    scores, box = top_box(models, image) if valid > 0 else ([0.0] * 3, None)
    (lat, lon), (slat, slon) = (np.ravel(v) for v in footprint_latlon(fit, [pas], [line], [side]))
    score = float(np.mean(scores))
    corners = [round(float(v), 1) for v in box] if box is not None else [""] * 4
    return {"group": group, "cycle": cycle, "pass": pas, "line": line, "side": side,
            "valid_frac": round(valid, 4),
            **{f"score_s{i}": round(s, 4) for i, s in enumerate(scores, 1)},
            "score": round(score, 4), "iw": int(score >= THRESHOLD),
            **dict(zip(("box_x0", "box_y0", "box_x1", "box_y1"), corners)),
            "lat": round(float(lat), 3), "lon": round(float(lon), 3),
            "swath_lat": round(float(slat), 3), "swath_lon": round(float(slon), 3),
            "labeled": int(group in labeled)}


def scan(pass_name: str, out: Path, limit: int | None, resolution: int) -> None:
    """Append one row per footprint of the pass to out, skipping footprints already written."""
    # pylint: disable=too-many-locals
    done = set()
    if out.exists():
        with open(out, encoding="utf-8") as fh:
            done = {row["group"] for row in csv.DictReader(fh)}
    models = [load_model(FINAL / name, resolution=resolution) for name in CKPTS]
    fit, labeled = load_fit(out.parent / "latlon_params.json"), labeled_groups()
    with h5py.File(H5, "r") as h5, open(out, "a", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, FIELDS, extrasaction="ignore")
        if not done:
            writer.writeheader()
        scenes = sorted(h5[pass_name])[:limit]
        tic = time.time()
        for n, scene in enumerate(scenes, 1):
            for side in ("L", "R"):
                group = f"{pass_name}/{scene}/{side}"
                if group not in done:
                    writer.writerow(footprint_row(h5, group, models, fit, labeled))
            if n % 200 == 0 or n == len(scenes):
                fh.flush()
                LOGGER.info("%s [%d/%d] scenes, %.2f s/scene", pass_name, n, len(scenes),
                            (time.time() - tic) / n)


def main() -> None:
    """Scan one pass; the pass comes from --pass or the Slurm array index."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--pass-name", required=True, help="h5 group, e.g. Pass_001")
    ap.add_argument("--out-dir", type=Path, default=Path("runs/scan_full"))
    ap.add_argument("--limit", type=int, default=None, help="first N scenes, for smoke tests")
    ap.add_argument("--resolution", type=int, default=560)
    opts = ap.parse_args()
    suffix = f"_limit{opts.limit}" if opts.limit else ""
    scan(opts.pass_name, opts.out_dir / f"{opts.pass_name}{suffix}.csv", opts.limit, opts.resolution)


if __name__ == "__main__":
    main()

def read_scan(scan_dir: Path) -> list[dict]:
    """All scan rows from the per-pass CSVs."""
    rows = []
    for path in sorted(glob.glob(str(scan_dir / "Pass_*.csv"))):
        with open(path, encoding="utf-8") as fh:
            rows.extend(csv.DictReader(fh))
    return rows