noaa-goes-visualization/filter_FITS.py

#!/usr/bin/env python
"""Index the SUVI archive and record per-frame quality verdicts.

Replaces the older version of this script, which encoded its verdict by renaming
files to ``_f``/``_e`` and wrote a diagnostic JPG next to every reject.  Verdicts now
go to the SQLite index, so the archive keeps NOAA's filenames, several detectors can
disagree about a frame, the continuous scores survive for later re-thresholding, and
nothing has to be re-read to change a threshold.

Subcommands::

    filter_FITS.py index     # sync the index with the archive
    filter_FITS.py scan      # cache FITS header metadata (~18 KB read per frame)
    filter_FITS.py detect    # run detectors and record verdicts
    filter_FITS.py plots     # regenerate diagnostic images for flagged frames

``index`` and ``scan`` are incremental: they skip what they already have, so they can
be re-run as the puller brings in new data.  ``index`` reads only the day-directories
whose mtime changed, which is what keeps re-indexing affordable -- see suvi/index.py
for why traversing this archive in full is something to avoid, not merely optimise.

A caution before using ``detect --detectors geometry_v1`` to gate anything: measured
over 20 sampled days of 2024, that method rejects 36.9% of frames, and does so
all-or-nothing -- 0% on most days, 73-100% on days with a large active region, almost
entirely for "failed centre".  Its centre test uses the intensity-weighted centroid,
which a bright active region displaces even when the header reports the disc exactly
centred.  Run bench.py before trusting any detector to discard data.
"""

import argparse
import json
import os
import sys
import time
from collections import defaultdict
from multiprocessing import Pool

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import db, detectors, fitsio, index, measures, metrics, paths


def parse_int_list(text):
    return tuple(int(part) for part in text.split(",") if part.strip())


# ---------------------------------------------------------------------------- index


def cmd_index(args):
    """Bring the index in step with the archive.

    Reads only the day-directories whose mtime changed since the last run, so a
    routine re-index costs ~7,400 directory stats rather than 2.65M file lookups.
    That ratio is not just a speed matter: a full traversal fills the guest's inode
    cache, and this archive lives on a virtiofs re-export of an SMB share whose daemon
    holds a host descriptor per cached inode.  Use --full only for a cold start or when
    the index is suspect.
    """
    root = args.root or paths.data_root()
    conn = db.connect(args.db)
    satellites = parse_int_list(args.satellites)
    wavelengths = parse_int_list(args.wavelengths)

    print(f"Indexing {root}" + (" (full rescan)" if args.full else " (changed directories only)"))
    summary = index.reconcile(
        conn,
        root,
        satellites,
        wavelengths,
        years=args.year,
        force=args.full,
        progress=print if args.progress else None,
    )
    total = conn.execute("SELECT count(*) c FROM frame").fetchone()["c"]
    print(
        f"  {summary['directories_changed']} of {summary['directories_checked']} "
        f"directories changed in {summary['seconds']:.0f}s"
    )
    print(f"  frames added {summary['frames_added']}, removed {summary['frames_removed']}")
    if summary["directories_vanished"]:
        print(f"  directories that disappeared: {summary['directories_vanished']}")
    dupes = summary.get("duplicate_slots") or []
    if dupes:
        print(f"  {len(dupes)} files share an observation slot with another and were "
              f"left out of the index (see --duplicates to list them)")
        if args.duplicates:
            for path in dupes:
                print(f"    {path}")
    print(f"  index now holds {total} frames")

    suffixed = conn.execute(
        "SELECT count(*) c FROM frame WHERE path LIKE '%\\_f.fits' ESCAPE '\\' "
        "   OR path LIKE '%\\_e.fits' ESCAPE '\\'"
    ).fetchone()["c"]
    if suffixed:
        print(f"\n{suffixed} indexed frames still carry _f/_e suffixes.")
        print("Run migrate_unrename.py to fold those verdicts into the index.")
    conn.close()
    return 0


# ----------------------------------------------------------------------------- scan


def _scan_one(job):
    frame_id, path = job
    values, error = fitsio.scan_header(path)
    return frame_id, values, error


def cmd_scan(args):
    fitsio.quiet_astropy()
    conn = db.connect(args.db)
    root = args.root or paths.data_root()
    pending = db.unscanned_frames(conn, args.limit, t_from=args.t_from)
    if not pending:
        print("Every indexed frame already has cached header metadata.")
        return 0

    print(f"Scanning headers for {len(pending)} frames with {args.workers} workers")
    jobs = [(row["id"], paths.abspath(row["path"], root)) for row in pending]
    started = time.time()
    done = failed = 0
    batch = []
    with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
        for frame_id, values, error in pool.imap_unordered(_scan_one, jobs, 32):
            batch.append((frame_id, values, error))
            done += 1
            failed += bool(error)
            if len(batch) >= 2000:
                db.record_headers(conn, batch)
                conn.commit()
                batch.clear()
                rate = done / max(time.time() - started, 1e-6)
                print(f"  {done}/{len(jobs)} ({rate:.0f}/s, {failed} unreadable)")
    if batch:
        db.record_headers(conn, batch)
        conn.commit()
    print(f"Scanned {done} frames in {time.time() - started:.0f}s; {failed} unreadable")
    return 0


# --------------------------------------------------------------------------- detect


def _detect_one(job):
    """Run the frame-level detectors on one file."""
    frame_id, slot, path, wanted = job
    verdicts = {}
    values, header_error = fitsio.scan_header(path)
    features = detectors.FrameFeatures(slot=slot, header=values, error=header_error)

    if "header_v1" in wanted:
        started = time.perf_counter_ns()
        verdict = detectors.header_v1(features)
        verdicts["header_v1"] = (verdict, (time.perf_counter_ns() - started) // 1000)

    pixel_detectors = [n for n in ("geometry_v1", "disc_v1") if n in wanted]
    if pixel_detectors:
        image, image_error = fitsio.read_image(path)
        for name in pixel_detectors:
            started = time.perf_counter_ns()
            if image is None:
                verdict = detectors.unreadable(image_error)
            elif name == "geometry_v1":
                verdict = detectors.geometry_v1(image, slot[1])
            else:
                verdict = detectors.disc_v1(image, slot[1], values)
            verdicts[name] = (verdict, (time.perf_counter_ns() - started) // 1000)

    return frame_id, verdicts


def cmd_detect(args):
    fitsio.quiet_astropy()
    conn = db.connect(args.db)
    root = args.root or paths.data_root()
    wanted = tuple(args.detectors.split(","))
    unknown = [name for name in wanted if name not in detectors.ALL_DETECTORS]
    if unknown:
        print(f"Unknown detectors: {unknown}.  Available: {list(detectors.ALL_DETECTORS)}")
        return 1

    frame_detectors = [n for n in wanted if n in detectors.FRAME_DETECTORS]
    series_detectors = [n for n in wanted if n in detectors.SERIES_DETECTORS]
    if series_detectors:
        print(f"NOTE: {series_detectors} need a whole series in memory; use bench.py "
              "for those, or run them here one band at a time via --wavelengths.")

    where, params = ["1=1"], []
    if args.start is not None:
        where.append("t_start >= ?")
        params.append(args.start)
    if args.end is not None:
        where.append("t_start < ?")
        params.append(args.end)
    if args.wavelengths:
        bands = parse_int_list(args.wavelengths)
        where.append(f"wavelength IN ({','.join('?' * len(bands))})")
        params.extend(bands)
    # header_v1 reads nothing but the header, and every header is already in the
    # index -- so a header-only run is a query, not a 5.8M-file traversal.  The values
    # are the same ones the file read would have produced (see db.header_from_row), so
    # this is a speed difference and not a judgement difference.
    from_cache = bool(frame_detectors) and set(frame_detectors) == {"header_v1"}
    if from_cache:
        columns = ", ".join(f"h.{name}" for name in db.HEADER_FIELDS)
        # `where` needs no qualifying: it filters on t_start and wavelength, and the
        # header table has neither column (its band is `wavelnth`), so the names are
        # unambiguous under the join.
        rows = conn.execute(
            f"SELECT f.id, f.satellite, f.wavelength, f.t_start, "
            f"{columns}, h.read_error FROM frame f "
            f"JOIN header h ON h.frame_id = f.id "
            f"WHERE {' AND '.join(where)} ORDER BY f.t_start",
            params,
        ).fetchall()
        if not rows:
            print("No indexed frames with a cached header match; run "
                  "`filter_FITS.py scan` first, or pass a pixel detector.")
            return 1
    else:
        rows = conn.execute(
            f"SELECT id, path, satellite, wavelength, t_start FROM frame "
            f"WHERE {' AND '.join(where)} ORDER BY t_start",
            params,
        ).fetchall()
        if not rows:
            print("No indexed frames match; run `filter_FITS.py index` first.")
            return 1

    runs = {
        name: db.create_detector_run(conn, name, {"driver": "filter_FITS"}, notes=args.notes)
        for name in frame_detectors
    }
    if from_cache:
        print(f"Running {frame_detectors} over {len(rows)} frames from the cached "
              f"headers (no files are opened)")
    else:
        print(f"Running {frame_detectors} over {len(rows)} frames "
              f"with {args.workers} workers")

    started = time.time()
    pending = defaultdict(list)
    tallies = {name: defaultdict(int) for name in frame_detectors}
    done = 0

    def collect(frame_id, verdicts):
        nonlocal done
        for name, (verdict, elapsed) in verdicts.items():
            pending[name].append(
                (frame_id, verdict.verdict, verdict.reason, verdict.scores, elapsed)
            )
            tallies[name][verdict.verdict] += 1
        done += 1
        if done % 5000 == 0:
            for name, batch in pending.items():
                db.record_detections(conn, runs[name], batch)
                batch.clear()
            conn.commit()
            print(f"  {done}/{len(rows)} ({done / max(time.time() - started, 1e-6):.0f}/s)")

    if from_cache:
        # No pool: judging a cached header is a few dictionary lookups, and handing
        # rows to worker processes would cost more in pickling than it saves.
        for row in rows:
            slot = (row["satellite"], row["wavelength"], row["t_start"])
            values, error = db.header_from_row(row)
            features = detectors.FrameFeatures(slot=slot, header=values, error=error)
            clock = time.perf_counter_ns()
            verdict = detectors.header_v1(features)
            collect(row["id"],
                    {"header_v1": (verdict, (time.perf_counter_ns() - clock) // 1000)})
    else:
        jobs = [
            (
                row["id"],
                (row["satellite"], row["wavelength"], row["t_start"]),
                paths.abspath(row["path"], root),
                frozenset(frame_detectors),
            )
            for row in rows
        ]
        with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
            for frame_id, verdicts in pool.imap_unordered(_detect_one, jobs, 16):
                collect(frame_id, verdicts)
    for name, batch in pending.items():
        if batch:
            db.record_detections(conn, runs[name], batch)
    conn.commit()

    print(f"\nDone in {time.time() - started:.0f}s")
    for name in frame_detectors:
        counts = tallies[name]
        total = sum(counts.values()) or 1
        print(f"  {name}: run {runs[name]} -- "
              + ", ".join(f"{v} {c} ({c / total:.1%})" for v, c in sorted(counts.items())))
    return 0


# ------------------------------------------------------------------------ calibrate

#: Absolute floor on each measurement's half-width, so a band whose sample happens to
#: be unusually tight does not get bounds narrower than the physics justifies.  Set
#: from the observed cross-band spread of each quantity.
CALIBRATION_FLOORS = {"limb_contrast": 0.15, "radius_ratio": 0.06, "limb_width": 0.06}

#: Which side of each measurement actually diagnoses a fault.  Bounding the other
#: side buys nothing and, on a wide-spread quantity, produces a negative limit that
#: can never trip -- which looks like a rule while being none.
#:
#: * limb_contrast falls towards zero as the disc disappears; it cannot exceed 1.
#: * radius_ratio is diagnostic in both directions -- too small or too large is wrong.
#: * limb_width grows as the limb smears; a sharper-than-usual limb is not a fault.
CALIBRATION_SIDES = {
    "limb_contrast": ("lower", 0.02, 1.0),
    "radius_ratio": ("both", None, None),
    "limb_width": ("upper", 0.0, None),
}


def _sample_frames(conn, root, wavelength, count, satellites=(16, 18)):
    """Pick frames spread evenly across the whole time range of one band.

    Queries the index rather than walking the archive -- which is the point of
    having an index, and avoids the traversal that fills the guest's inode cache
    (see suvi/index.py).

    Deliberately does *not* filter on the legacy ``_f``/``_e`` verdicts.  Those were
    written in May 2024 by a filter retuned that July, so selecting on them would
    calibrate the new detector against the old one's mistakes -- including the very
    active-region rejections this detector exists to stop.
    """
    placeholders = ",".join("?" * len(satellites))
    rows = conn.execute(
        f"SELECT path FROM frame WHERE wavelength = ? AND satellite IN ({placeholders}) "
        f"ORDER BY t_start",
        (wavelength, *satellites),
    ).fetchall()
    if not rows:
        return []
    stride = max(1, len(rows) // count)
    return [paths.abspath(row["path"], root) for row in rows[::stride]][:count]


def _calibration_measure(path):
    """Measure one frame, skipping any the header already condemns."""
    values, error = fitsio.scan_header(path)
    if error or not values:
        return None
    if values.get("eclipse") or values.get("empty"):
        return None  # genuinely bad; must not widen the good-frame bounds
    diameter = values.get("diam_sun")
    mean = values.get("img_mean")
    if not diameter or mean is None:
        return None
    low, high = detectors.HEADER_MEAN_BOUNDS.get(
        int(values.get("wavelnth") or 0), (0.0, float("inf"))
    )
    if not (low <= mean <= high):
        return None  # dropout or blowout
    image, image_error = fitsio.read_image(path)
    if image is None:
        return None
    return detectors.disc_profile(image, float(diameter) / 4.0)


def cmd_calibrate(args):
    fitsio.quiet_astropy()
    root = args.root or paths.data_root()
    conn = db.connect(args.db, readonly=True)
    wavelengths = parse_int_list(args.wavelengths)
    margin = args.margin
    table = {}

    print(f"Calibrating disc_v1 from {root}")
    print(f"  {args.samples_per_band} frames per band, margin {margin}x the observed span\n")

    for wavelength in wavelengths:
        candidates = _sample_frames(conn, root, wavelength, args.samples_per_band * 2)
        if not candidates:
            print(f"{wavelength:>5}: no frames found")
            continue
        results = []
        with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
            for measured in pool.imap_unordered(_calibration_measure, candidates, 8):
                if measured is None:
                    continue
                if any(v is None for v in measured.values()):
                    results.append(measured)  # keep, so unmeasurable rate is visible
                    continue
                results.append(measured)
                if len(results) >= args.samples_per_band:
                    break

        usable = [r for r in results if all(v is not None for v in r.values())]
        unmeasurable = len(results) - len(usable)
        if len(usable) < 20:
            print(f"{wavelength:>5}: only {len(usable)} usable frames; skipping")
            continue

        bounds = {}
        print(f"{wavelength:>5}A  n={len(usable)}  unmeasurable={unmeasurable}")
        for name in ("limb_contrast", "radius_ratio", "limb_width"):
            values = np.array([r[name] for r in usable], dtype=float)
            low, high = float(values.min()), float(values.max())
            span = high - low
            half = max(margin * span, CALIBRATION_FLOORS[name])
            side, floor, ceiling = CALIBRATION_SIDES[name]

            lower = low - half if side in ("lower", "both") else floor
            upper = high + half if side in ("upper", "both") else ceiling
            # Keep each limit on the physical side of zero.  A negative floor on a
            # non-negative quantity is not a loose rule, it is no rule at all.
            if floor is not None:
                lower = max(lower, floor)
            if ceiling is not None:
                upper = min(upper, ceiling)
            bounds[name] = (round(lower, 4), round(upper, 4))
            outside = int(((values < lower) | (values > upper)).sum())
            print(f"        {name:<14} observed {low:.4f}-{high:.4f}  "
                  f"p50 {np.median(values):.4f}  ->  bounds {lower:.4f}-{upper:.4f}  "
                  f"sample failures {outside}/{len(values)}")
        table[wavelength] = bounds

    if not table:
        print("\nNothing calibrated.")
        return 1

    print("\n" + "=" * 78)
    print("Paste into DISC_BOUNDS in suvi/detectors.py:\n")
    print("DISC_BOUNDS = {")
    for wavelength, bounds in sorted(table.items()):
        print(f"    {wavelength}: {{")
        for name, (lower, upper) in bounds.items():
            print(f'        "{name}": ({lower}, {upper}),')
        print("    },")
    print("}")
    if args.json:
        with open(args.json, "w") as handle:
            json.dump(table, handle, indent=2)
        print(f"\nAlso written to {args.json}")
    return 0


# --------------------------------------------------------------------------- sample

#: The production combination: a frame is bad if either detector says so.
SAMPLE_DETECTORS = ("header_v1", "disc_v1")


def _judge_timestamp(job):
    """Read all six bands for one (satellite, timestamp) and judge each.

    Returns (satellite, timestamp, {band: (verdict, reason, scores)}), or None if
    the set is incomplete -- a composite needs all six.
    """
    satellite, timestamp, band_paths = job
    verdicts = {}
    for band, path in sorted(band_paths.items()):
        values, header_error = fitsio.scan_header(path)
        features = detectors.FrameFeatures(
            slot=(satellite, band, timestamp), header=values, error=header_error
        )
        header = detectors.header_v1(features)
        image, image_error = fitsio.read_image(path)
        if image is None:
            disc = detectors.unreadable(image_error)
        else:
            disc = detectors.disc_v1(image, band, values)

        # metrics.combine_verdicts, not a hand-rolled fold: the inline version
        # silently counted UNKNOWN as "good", which the shared policy refuses.
        combined = metrics.combine_verdicts(
            [{band: header}, {band: disc}], policy="any"
        )[band]
        verdicts[band] = (combined.verdict, combined.reason,
                         {**header.scores, **disc.scores})
    return satellite, timestamp, verdicts


def _render_sample(job):
    """Render one annotated composite for review."""
    import merger_FITS
    from PIL import ImageDraw, ImageFont

    satellite, timestamp, band_paths, caption, out_path = job
    try:
        arrays = []
        for band in sorted(band_paths):
            image, error = fitsio.read_image(band_paths[band])
            if image is None:
                return out_path, f"band {band}: {error}"
            arrays.append(image)
        picture = merger_FITS.composite_from_arrays(arrays, timestamp)

        font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                 "OpenSans-Regular.ttf")
        font = ImageFont.truetype(font_path, size=20)
        draw = ImageDraw.Draw(picture)
        # Caption along the bottom, where it covers no solar structure.
        y = picture.size[1] - 26 * len(caption)
        for line in caption:
            draw.text((12, y), line, (255, 235, 120), font=font)
            y += 26
        picture.save(out_path, quality=95)
        return out_path, None
    except Exception as exc:
        return out_path, f"{type(exc).__name__}: {exc}"


def cmd_sample(args):
    """Write folders of suspected-good and suspected-bad frames for manual review.

    The bench measures detectors against damage it injected itself, which proves
    they find what was planted but says nothing about what they do to real frames
    nobody labelled.  This is that check: a person looks at what the detector kept
    and what it threw away, and says whether they agree.

    Each image is the full composite for the frame's timestamp and satellite -- which
    already tiles all six bands down its sides -- captioned with the verdict, the
    band responsible, and the scores behind it.
    """
    fitsio.quiet_astropy()
    root = args.root or paths.data_root()
    conn = db.connect(args.db, readonly=True)
    satellites = parse_int_list(args.satellites)

    # Candidate timestamps spread evenly across whatever the index holds, so the
    # sample is not a picture of one week.
    rows = conn.execute(
        f"SELECT satellite, t_start, wavelength, path FROM frame "
        f"WHERE satellite IN ({','.join('?' * len(satellites))}) "
        f"ORDER BY t_start",
        satellites,
    ).fetchall()
    conn.close()

    by_key = defaultdict(dict)
    for row in rows:
        by_key[(row["satellite"], row["t_start"])][row["wavelength"]] = paths.abspath(
            row["path"], root
        )
    complete = [k for k, v in by_key.items() if len(v) == 6]
    if not complete:
        print("No timestamp has all six bands; cannot build composites.")
        return 1

    stride = max(1, len(complete) // max(args.scan, 1))
    candidates = sorted(complete)[::stride][: args.scan]
    print(f"{len(complete)} complete timestamps in the index; judging {len(candidates)}")

    good, bad = [], []
    jobs = [(sat, when, by_key[(sat, when)]) for sat, when in candidates]

    with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
        for result in pool.imap_unordered(_judge_timestamp, jobs, 4):
            if result is None:
                continue
            satellite, timestamp, verdicts = result
            flagged = {b: v for b, v in verdicts.items() if v[0] == "bad"}
            if flagged and len(bad) < args.n:
                bad.append((satellite, timestamp, flagged))
            elif not flagged and len(good) < args.n:
                good.append((satellite, timestamp, verdicts))
            if len(good) >= args.n and len(bad) >= args.n:
                break

    print(f"  collected {len(bad)} flagged and {len(good)} clean timestamps")
    if len(bad) < args.n:
        print(f"  NOTE: only {len(bad)} flagged found in {len(candidates)} scanned; "
              f"raise --scan for more.")

    render_jobs = []
    for folder, items in (("bad", bad), ("good", good)):
        out_dir = os.path.join(args.out, folder)
        os.makedirs(out_dir, exist_ok=True)
        for satellite, timestamp, verdicts in items:
            stamp = time.strftime("%Y%m%dT%H%M%S", time.gmtime(timestamp))
            if folder == "bad":
                bands = ",".join(f"{b}A" for b in sorted(verdicts))
                first = verdicts[sorted(verdicts)[0]]
                # Lead the filename with the reason so the folder sorts by failure mode.
                slug = "".join(
                    ch if ch.isalnum() else "_" for ch in (first[1] or "flagged")
                )[:40]
                caption = [
                    f"FLAGGED  g{satellite}  {stamp}  bands: {bands}",
                    (first[1] or "")[:110],
                    " ".join(f"{k}={v:.3g}" for k, v in list(first[2].items())[:5]),
                ]
                name = f"{slug}__g{satellite}_{stamp}.jpg"
            else:
                sample_scores = verdicts[sorted(verdicts)[0]][2]
                caption = [
                    f"CLEAN  g{satellite}  {stamp}  all six bands passed",
                    " ".join(f"{k}={v:.3g}" for k, v in list(sample_scores.items())[:5]),
                ]
                name = f"g{satellite}_{stamp}.jpg"
            render_jobs.append(
                (satellite, timestamp, by_key[(satellite, timestamp)], caption,
                 os.path.join(out_dir, name))
            )

    print(f"Rendering {len(render_jobs)} composites...")
    failures = 0
    with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
        for path, error in pool.imap_unordered(_render_sample, render_jobs, 2):
            if error:
                failures += 1
                if failures <= 3:
                    print(f"  failed {os.path.basename(path)}: {error}")
    print(f"  wrote {len(render_jobs) - failures} images to {args.out}")
    print(f"  review {args.out}/bad  -- nothing here should be a frame you would keep")
    print(f"  review {args.out}/good -- nothing here should be a frame you would discard")
    return 0


# ---------------------------------------------------------------------------- plots


def cmd_plots(args):
    """Regenerate the diagnostic image the old filter wrote beside every reject.

    Kept as an on-demand tool rather than a side effect of filtering: the plots are
    derivable from the index plus the FITS, and writing millions of them into the
    archive was what made the old approach hard to undo.
    """
    import matplotlib

    matplotlib.use("Agg")
    import numpy as np
    from matplotlib import pyplot as plt

    fitsio.quiet_astropy()
    conn = db.connect(args.db, readonly=True)
    root = args.root or paths.data_root()
    run_id = db.latest_run_id(conn, args.detector)
    if run_id is None:
        print(f"No run named {args.detector!r}")
        return 1

    rows = conn.execute(
        """
        SELECT f.path, f.wavelength, d.reason FROM detection d
        JOIN frame f ON f.id = d.frame_id
        WHERE d.run_id = ? AND d.verdict = 'bad' ORDER BY f.t_start LIMIT ?
        """,
        (run_id, args.limit),
    ).fetchall()
    os.makedirs(args.out, exist_ok=True)
    print(f"Writing {len(rows)} diagnostic plots to {args.out}")

    reference = detectors.ideal_disc_axis()
    axis = list(range(detectors.HALF_DIMS))
    for row in rows:
        path = paths.abspath(row["path"], root)
        image, error = fitsio.read_image(path)
        if image is None:
            continue
        import cv2 as cv

        threshold = detectors.GEOMETRY_THRESHOLDS[row["wavelength"]]
        data = cv.resize(
            np.nan_to_num(image),
            (detectors.HALF_DIMS, detectors.HALF_DIMS),
            interpolation=cv.INTER_LINEAR,
        )
        shown = np.clip(data, None, threshold) / threshold
        xavg, yavg = np.average(shown, 0), np.average(shown, 1)

        figure, axes = plt.subplots(
            2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1),
            gridspec_kw={"hspace": 0.0, "wspace": 0.0}, figsize=(8, 8),
        )
        axes[0][0].axis("off")
        axes[0][1].plot(axis, xavg)
        axes[0][1].plot(axis, reference)
        axes[1][0].plot(yavg, axis)
        axes[1][0].plot(reference, axis)
        axes[1][0].invert_xaxis()
        axes[1][0].invert_yaxis()
        axes[1][1].imshow(shown, aspect="auto", vmin=0, vmax=1)
        axes[0][1].sharex(axes[1][1])
        axes[1][0].sharey(axes[1][1])
        figure.suptitle(f"{os.path.basename(row['path'])}\n{row['reason']}", fontsize=8)
        figure.savefig(
            os.path.join(args.out, os.path.basename(row["path"]).replace(".fits", ".jpg")),
            dpi=110,
        )
        plt.close(figure)
    return 0


# ------------------------------------------------------------------------------ cli


def _day(t):
    return time.strftime("%Y-%m-%d", time.gmtime(t))


# -------------------------------------------------------------------------- measure

MEASURE_RUN = detectors.MEASURE_RUN
#: One measuring job is one (satellite, band) stream for this long, plus a margin
#: either side for the structure score.
MEASURE_DAY = 86400
#: Jobs are built and run a week at a time.
MEASURE_CHUNK = 7 * MEASURE_DAY
#: Stream-days (~400 frames each) a worker measures before it is replaced.  The
#: archive-wide judgment's workers grew to ~3.6 GB each over 20 hours and, with four
#: of them, took the machine into an OOM; a fresh process every ten thousand frames
#: bounds whatever a worker accumulates, at the cost of a second's start-up.
MEASURE_RECYCLE = 25


#: Attempts at a read the file system refused, pausing READ_BACKOFF seconds and
#: doubling (~2 minutes in all), before the pass stops rather than store a verdict
#: about a frame it never saw.
READ_ATTEMPTS = 8
READ_BACKOFF = 1.0


def _retrying(read, path):
    for attempt in range(READ_ATTEMPTS):
        try:
            return read(path)
        except fitsio.TransientReadError:
            if attempt == READ_ATTEMPTS - 1:
                raise
            time.sleep(READ_BACKOFF * 2 ** attempt)


def measure_band(frames):
    """Every per-frame measurement for one (satellite, band) stretch.

    `frames` is [(frame_id, slot, path)] in time order, including a margin either
    side for the structure score.  Each frame gets the single-frame judgment
    (``detectors.judge_frame``) and every measure of ``suvi.measures``; then the
    structure score, which compares thumbnails that are not kept, against the
    nearest frames that passed.  Returns {slot: (frame_id, source, verdict,
    elapsed_us)}, the verdict the single-frame one with every measure in its scores.
    """
    series, decided, good = [], {}, set()
    for frame_id, slot, path in frames:
        started = time.perf_counter_ns()
        cards, header_error = _retrying(fitsio.read_header, path)
        image_box = {}

        def read(path=path, box=image_box):
            box["image"], box["error"] = _retrying(fitsio.read_image, path)
            return box["image"], box["error"]

        source, verdict, image = detectors.judge_frame(slot, cards, header_error, read)
        values = fitsio.header_metadata(cards)
        scores = {**verdict.scores, **measures.header_measures(cards)}
        if image is not None:
            scores.update(measures.pixel_measures(image, values.get("diam_sun")))
        if verdict.verdict != detectors.BAD:
            good.add(slot)
        error = header_error or (image_box.get("error") if image is None and image_box else None)
        series.append(detectors.FrameFeatures(
            slot=slot, header=values, error=error,
            thumbnail=detectors.thumbnail(image) if image is not None else None))
        decided[slot] = (frame_id, source, detectors.Verdict(verdict.verdict, verdict.reason, scores),
                         (time.perf_counter_ns() - started) // 1000)

    for features, temporal in zip(series, detectors.temporal_v2(series, good)):
        frame_id, source, verdict, elapsed = decided[features.slot]
        scores = {**verdict.scores,
                  **{f"temporal_{name}": value for name, value in temporal.scores.items()}}
        decided[features.slot] = (frame_id, source,
                                  detectors.Verdict(verdict.verdict, verdict.reason, scores),
                                  elapsed)
    return decided


def _measure_job(job):
    frames, wanted = job
    out = []
    for frame_id, source, verdict, elapsed in measure_band(frames).values():
        if frame_id not in wanted:
            continue
        reason = f"{source}: {verdict.reason}" if source and verdict.reason else verdict.reason
        out.append((frame_id, verdict.verdict, reason, verdict.scores, elapsed))
    return out


def measure_jobs(conn, root, start, end, satellites, run_id):
    """One job per (satellite, band, day) that still has unmeasured frames."""
    margin = detectors.TEMPORAL2_REACH * paths.CADENCE
    jobs = []
    for day in range(start - start % MEASURE_DAY, end, MEASURE_DAY):
        lo, hi = max(day, start), min(day + MEASURE_DAY, end)
        for satellite in satellites:
            for wavelength in paths.WAVELENGTHS:
                rows = conn.execute(
                    "SELECT f.id, f.path, f.t_start, EXISTS (SELECT 1 FROM detection d "
                    "WHERE d.run_id = ? AND d.frame_id = f.id) AS done FROM frame f "
                    "WHERE f.satellite = ? AND f.wavelength = ? AND f.t_start >= ? "
                    "AND f.t_start < ? ORDER BY f.t_start",
                    (run_id, satellite, wavelength, lo - margin, hi + margin)).fetchall()
                wanted = [r["id"] for r in rows if lo <= r["t_start"] < hi and not r["done"]]
                if wanted:
                    jobs.append(([(r["id"], (satellite, wavelength, r["t_start"]),
                                   paths.abspath(r["path"], root)) for r in rows], wanted))
    return jobs


def cmd_measure(args):
    """Store every frame's measurements for a time range, as one detector run.

    The expensive half of judging the archive -- reading and decoding every frame
    -- done once; ``judge`` then works from what this stores.  Resumable: rerun
    the same command and stream-days already measured under the same
    configuration are skipped.
    """
    fitsio.quiet_astropy()
    conn = db.connect(args.db)
    root = args.root or paths.data_root()
    config = detectors.measure_config()
    run_id = None if args.new_run else db.latest_run_id(conn, MEASURE_RUN, config)
    if run_id is None:
        run_id = db.create_detector_run(conn, MEASURE_RUN, config, notes=args.notes)
        print(f"new measure run {run_id}")
    else:
        print(f"resuming measure run {run_id}")
    satellites = parse_int_list(args.satellites)
    marks = ",".join("?" * len(satellites))
    total = conn.execute(
        f"SELECT COUNT(*) FROM frame f WHERE f.t_start >= ? AND f.t_start < ? "
        f"AND f.satellite IN ({marks}) AND NOT EXISTS (SELECT 1 FROM detection d "
        f"WHERE d.run_id = ? AND d.frame_id = f.id)",
        (args.start, args.end, *satellites, run_id)).fetchone()[0]
    print(f"{total} frames to measure with {args.workers} workers", flush=True)
    started, done = time.time(), 0
    pool = Pool(args.workers, initializer=fitsio.quiet_astropy,
                maxtasksperchild=MEASURE_RECYCLE)
    try:
        for chunk in range(args.start, args.end, MEASURE_CHUNK):
            jobs = measure_jobs(conn, root, chunk, min(chunk + MEASURE_CHUNK, args.end),
                                satellites, run_id)
            for results in pool.imap_unordered(_measure_job, jobs):
                db.record_detections(conn, run_id, results)
                conn.commit()                  # a stream-day is the unit of resume
                done += len(results)
            rate = done / max(time.time() - started, 1e-6)
            print(f"  {_day(chunk)}: {done}/{total} ({rate:.1f}/s, "
                  f"{(total - done) / max(rate, 1e-6) / 3600:.1f} h left)", flush=True)
    finally:
        # Explicit join: `with Pool` only terminates, and an unreaped result
        # handler is how the corpus builder once grew to 16.6 GB.
        pool.terminate()
        pool.join()
    print(f"done in {(time.time() - started) / 3600:.2f} h; run {run_id}")
    return 0


# ---------------------------------------------------------------------------- judge

JUDGE_RUN = detectors.JUDGMENT_RUN
#: Frames loaded either side of a judged stretch, in slots: the baseline reaches
#: two hours out, and whether those references were flagged depends on theirs.
JUDGE_MARGIN = 2 * detectors.BASELINE_REACH
#: Judged a week at a time, one satellite at a time, to bound memory.
JUDGE_CHUNK = 7 * 86400


def load_measured(conn, run_id, satellite, start, end):
    """One satellite's stored measurements for [start, end):
    {band: [(frame_id, FrameFeatures, source, Verdict)]} in time order, the
    Verdict the stored single-frame one with every measure in its scores."""
    streams = defaultdict(list)
    for row in conn.execute(
            "SELECT f.id, f.wavelength, f.t_start, d.verdict, d.reason, d.scores_json, "
            "h.img_mean FROM frame f JOIN detection d ON d.frame_id = f.id AND d.run_id = ? "
            "LEFT JOIN header h ON h.frame_id = f.id WHERE f.satellite = ? "
            "AND f.t_start >= ? AND f.t_start < ? ORDER BY f.wavelength, f.t_start",
            (run_id, satellite, start, end)):
        scores = json.loads(row["scores_json"]) if row["scores_json"] else {}
        source, reason = None, row["reason"]
        if row["verdict"] == detectors.BAD and reason and ": " in reason:
            source, reason = reason.split(": ", 1)
        slot = (satellite, row["wavelength"], row["t_start"])
        features = detectors.FrameFeatures(slot=slot, header={"img_mean": row["img_mean"]},
                                           scores=scores)
        streams[row["wavelength"]].append(
            (row["id"], features, source, detectors.Verdict(row["verdict"], reason, scores)))
    return streams


def judge_measured(streams):
    """The production judgment from stored measurements.

    Per band: the stored single-frame verdict, then ``temporal_v2`` (from its
    stored scores) on each frame, then ``relative_v2`` against the frames that
    passed the single-frame checks; then ``pointing_v1`` reads what the
    satellite's other bands failed in each slot.  `streams` is as from
    ``load_measured``.  Returns {frame_id: (slot, source, Verdict)}, `source`
    naming the check that decided a BAD verdict.
    """
    decided = {}
    for band, frames in streams.items():
        single, good = {}, set()
        for frame_id, features, source, verdict in frames:
            if verdict.verdict != detectors.BAD:
                good.add(features.slot)
                temporal = detectors.temporal_verdict(features.scores)
                if temporal.verdict == detectors.BAD:
                    source, verdict = "temporal_v2", detectors.Verdict(
                        detectors.BAD, temporal.reason, features.scores)
            single[features.slot] = (frame_id, source, verdict)
        series = [features for _, features, _, _ in frames]
        for features, relative in zip(series, detectors.relative_v2(series, good)):
            frame_id, source, verdict = single[features.slot]
            scores = {**verdict.scores, **relative.scores}
            if verdict.verdict != detectors.BAD and relative.verdict == detectors.BAD:
                source, verdict = "relative_v2", detectors.Verdict(
                    detectors.BAD, relative.reason, scores)
            else:
                verdict = detectors.Verdict(verdict.verdict, verdict.reason, scores)
            decided[frame_id] = (features.slot, source, verdict)

    failed = defaultdict(dict)
    for slot, source, verdict in decided.values():
        failed[slot[2]][slot[1]] = source if verdict.verdict == detectors.BAD else None
    for frame_id, (slot, source, verdict) in list(decided.items()):
        if verdict.verdict != detectors.BAD:
            pointing = detectors.pointing_v1(slot[1], failed[slot[2]])
            if pointing.verdict == detectors.BAD:
                decided[frame_id] = (slot, "pointing_v1", detectors.Verdict(
                    detectors.BAD, pointing.reason, {**verdict.scores, **pointing.scores}))
    return decided


def cmd_judge(args):
    """Store the production judgment for a time range, as one detector run.

    Made from the stored measurements (``measure``), so it reads no frame and runs
    in minutes: a changed rule or limit is a new judgment, not a new pass over the
    archive.  Every consumer that decides what counts as bad -- archive profiling,
    the merger, the corpus builder's availability trace, the real-data renderer --
    reads this one run.  Resumable: frames already judged in the run are skipped.
    """
    conn = db.connect(args.db)
    measure_run = args.measure_run or detectors.measure_run(conn)
    if measure_run is None:
        raise SystemExit("No measurements under the current rules; run "
                         "`filter_FITS.py measure` over the range first")
    satellites = parse_int_list(args.satellites)
    marks = ",".join("?" * len(satellites))
    margin = JUDGE_MARGIN * paths.CADENCE
    unmeasured = conn.execute(
        f"SELECT COUNT(*) FROM frame f WHERE f.t_start >= ? AND f.t_start < ? "
        f"AND f.satellite IN ({marks}) AND NOT EXISTS (SELECT 1 FROM detection d "
        f"WHERE d.run_id = ? AND d.frame_id = f.id)",
        (args.start - margin, args.end + margin, *satellites, measure_run)).fetchone()[0]
    if unmeasured:
        raise SystemExit(f"{unmeasured} frames in the range (with its {JUDGE_MARGIN}-slot "
                         f"margin) have no measurements in run {measure_run}; "
                         "measure them first")
    config = detectors.judgment_config()
    run_id = None if args.new_run else db.latest_run_id(conn, JUDGE_RUN, config)
    if run_id is None:
        run_id = db.create_detector_run(conn, JUDGE_RUN, config, notes=args.notes)
        print(f"new judge run {run_id} from measure run {measure_run}")
    else:
        print(f"resuming judge run {run_id} from measure run {measure_run}")
    started, done, tallies = time.time(), 0, defaultdict(int)
    for chunk in range(args.start, args.end, JUDGE_CHUNK):
        hi = min(chunk + JUDGE_CHUNK, args.end)
        for satellite in satellites:
            judged = {r[0] for r in conn.execute(
                "SELECT d.frame_id FROM detection d JOIN frame f ON f.id = d.frame_id "
                "WHERE d.run_id = ? AND f.satellite = ? AND f.t_start >= ? AND f.t_start < ?",
                (run_id, satellite, chunk, hi))}
            decided = judge_measured(load_measured(conn, measure_run, satellite,
                                                   chunk - margin, hi + margin))
            rows = []
            for frame_id, (slot, source, verdict) in decided.items():
                if not chunk <= slot[2] < hi or frame_id in judged:
                    continue
                reason = (f"{source}: {verdict.reason}" if source and verdict.reason
                          else verdict.reason)
                rows.append((frame_id, verdict.verdict, reason, verdict.scores, 0))
                tallies[(verdict.verdict, source or "")] += 1
            db.record_detections(conn, run_id, rows)
            conn.commit()
            done += len(rows)
        print(f"  {_day(chunk)}: {done} judged", flush=True)
    print(f"done in {(time.time() - started) / 60:.1f} min; run {run_id}")
    for (verdict, source), count in sorted(tallies.items()):
        print(f"  {verdict:<8}{source or '-':<14}{count}")
    return 0


def build_parser():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)")
    parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)")
    subparsers = parser.add_subparsers(dest="command", required=True)

    index_cmd = subparsers.add_parser(
        "index", help="sync the index with the archive (only changed directories)"
    )
    index_cmd.add_argument("--satellites", default="16,18,19")
    index_cmd.add_argument("--wavelengths", default="94,131,171,195,284,304")
    index_cmd.add_argument("--year", type=int, action="append", default=None)
    index_cmd.add_argument("--full", action="store_true",
                           help="re-read every directory, ignoring recorded mtimes")
    index_cmd.add_argument("--progress", action="store_true")
    index_cmd.add_argument("--duplicates", action="store_true",
                           help="list files that share an observation slot")
    index_cmd.set_defaults(func=cmd_index)

    scan = subparsers.add_parser("scan", help="cache FITS header metadata")
    scan.add_argument("--workers", type=int, default=16)
    scan.add_argument("--limit", type=int, default=None)
    scan.add_argument("--t-from", type=int, default=None,
                      help="only frames at or after this unix time (era priority)")
    scan.set_defaults(func=cmd_scan)

    detect = subparsers.add_parser("detect", help="run detectors and record verdicts")
    detect.add_argument("--detectors", default="header_v1")
    detect.add_argument("--wavelengths", default=None)
    detect.add_argument("--start", type=int, default=None, help="Unix time, inclusive")
    detect.add_argument("--end", type=int, default=None, help="Unix time, exclusive")
    detect.add_argument("--workers", type=int, default=16)
    detect.add_argument("--notes", default=None)
    detect.set_defaults(func=cmd_detect)

    measure = subparsers.add_parser(
        "measure", help="store every frame's measurements for a time range (resumable)")
    measure.add_argument("--start", type=int, required=True, help="Unix time, inclusive")
    measure.add_argument("--end", type=int, required=True, help="Unix time, exclusive")
    measure.add_argument("--satellites", default="16,18")
    measure.add_argument("--workers", type=int, default=4)
    measure.add_argument("--new-run", action="store_true",
                         help="start a new run even if one with this configuration exists")
    measure.add_argument("--notes", default=None)
    measure.set_defaults(func=cmd_measure)

    judge = subparsers.add_parser(
        "judge", help="store the production judgment for a time range (resumable)")
    judge.add_argument("--start", type=int, required=True, help="Unix time, inclusive")
    judge.add_argument("--end", type=int, required=True, help="Unix time, exclusive")
    judge.add_argument("--satellites", default="16,18")
    judge.add_argument("--measure-run", type=int, default=None,
                       help="measurement run id (default: the one under the current rules)")
    judge.add_argument("--new-run", action="store_true",
                       help="start a new run even if one with this configuration exists")
    judge.add_argument("--notes", default=None)
    judge.set_defaults(func=cmd_judge)

    calibrate = subparsers.add_parser(
        "calibrate", help="derive per-band disc_v1 bounds from a sample of good frames"
    )
    calibrate.add_argument("--wavelengths", default="94,131,171,195,284,304")
    calibrate.add_argument("--samples-per-band", type=int, default=300)
    calibrate.add_argument("--margin", type=float, default=1.0,
                           help="half-width added either side, in multiples of the observed span")
    calibrate.add_argument("--workers", type=int, default=16)
    calibrate.add_argument("--json", default=None)
    calibrate.set_defaults(func=cmd_calibrate)

    sample = subparsers.add_parser(
        "sample", help="write good/bad frame folders for manual review"
    )
    sample.add_argument("--n", type=int, default=100, help="frames per folder")
    sample.add_argument("--out", required=True)
    sample.add_argument("--scan", type=int, default=4000,
                        help="timestamps to judge while looking for N of each")
    sample.add_argument("--satellites", default="16,18")
    sample.add_argument("--workers", type=int, default=8)
    sample.set_defaults(func=cmd_sample)

    plots = subparsers.add_parser("plots", help="regenerate diagnostics for flagged frames")
    plots.add_argument("--detector", default="geometry_v1")
    plots.add_argument("--out", default="diagnostics")
    plots.add_argument("--limit", type=int, default=200)
    plots.set_defaults(func=cmd_plots)
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    return args.func(args) or 0


if __name__ == "__main__":
    sys.exit(main())