#!/usr/bin/env python
"""Render a filler over real archive data: no synthetic damage, no truth.

    SUVI_MODEL=.../v8_best.pt render_real.py --start 2024-01-01 --end 2025-01-01 \\
        --out /media/shared/NOAA_GOES/renders/real2024_v8

The bench measures the model against damage it injected and truth it kept.  This
is the other thing: the archive as it really is, where the only repairs are for
frames the archive itself lost or spoiled, and nothing exists to score them
against.  The output is the product -- one composite per 4-minute slot, with bad
bands reconstructed by the model -- not a number.

What counts as bad is the **stored production judgment** (``filter_FITS.py
judge``: header, lunar, disc, field and clean-neighbour checks), the one run every
consumer reads, scored on hand-confirmed real defects.  An earlier version judged
frames here with header_v1 + disc_v1 alone; that pair passed 41% of the real
defects in the audit -- bars, ghost discs, static -- and the renderer composited
them, which is what the first 2024 video showed.  A day whose frames the run has
not judged is refused rather than rendered on a weaker test.

``--filler optical_flow`` repairs each bad band from its nearest clean frames
either side (the classical repair, the same as ``bench.py render --variant new``)
instead of the model: the baseline a learned filler is compared against on real
data.

Work proceeds a day at a time.  Each day is seen through a window reaching two
days back and two forward, which is enough for the model's +/-128-slot stack and
guarantees that any gap overlapping the day is seen at its true length up to
the 24-hour fill policy, so a long outage is left black rather than invented.
Gaps are judged band by band: a slot is only out of policy when some band has
gone unobserved by *both* satellites for longer than that.  GOES-16 lacking
94/131 while GOES-18 lacks 304 is not an outage -- every band is observed, and
the model takes each band from whichever frames hold it.  Every slot on the
cadence timeline is composited, not just slots the archive holds: a short outage
on both satellites is still in policy and still gets filled.

Resumable: a day is marked done only once all its composites are written, and
composites are never rewritten.
"""

import argparse
import datetime
import json
import os
import sys
import time

import numpy as np

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

import bench
from suvi import availability, cases, dataset, db, detectors, fillers, fitsio, paths

DAY = 86400
#: Days of context either side of the day being rendered.  Must cover the
#: model's stack (+/-128 slots is ~8.5 h) and the dual-gap policy (24 h).
BEFORE_DAYS = 2
AFTER_DAYS = 2


def day_starts(start, end):
    """UTC midnights in [start, end)."""
    return list(range(start - start % DAY, end, DAY))


def window_for(day):
    """The context window a day is rendered through, as [t0, t1)."""
    return day - BEFORE_DAYS * DAY, day + (1 + AFTER_DAYS) * DAY


def day_name(day):
    return datetime.datetime.fromtimestamp(day, datetime.timezone.utc).strftime("%Y-%m-%d")


#: Fillers a render can use; each writes ``<out>/<filler>/goes<satellite>``.
FILLERS = ("learned", "optical_flow")
#: Anchor frames the optical-flow filler keeps between fills.  Consecutive bad
#: slots share anchors; the bound keeps a bar-storm day from holding hundreds.
FLOW_CACHE = 64


class State:
    """Resume state under ``<out>/_state``: which (day, satellite) are finished
    for each filler, so two fillers rendered into one root never share markers."""

    def __init__(self, out, filler="learned"):
        self.done = os.path.join(out, "_state", "done", filler)
        os.makedirs(self.done, exist_ok=True)

    def is_done(self, day, satellite):
        return os.path.exists(os.path.join(self.done, f"{day_name(day)}_g{satellite}"))

    def mark_done(self, day, satellite, summary):
        with open(os.path.join(self.done, f"{day_name(day)}_g{satellite}"), "w") as handle:
            json.dump(summary, handle)


def frame_rows(conn, t0, t1, satellites, run_id):
    """Index rows for [t0, t1) with the run's verdict:
    (satellite, wavelength, t_start, path, size_bytes, verdict or None)."""
    marks = ",".join("?" * len(satellites))
    return [
        (row["satellite"], row["wavelength"], row["t_start"], row["path"],
         row["size_bytes"], row["verdict"])
        for row in conn.execute(
            f"SELECT f.satellite, f.wavelength, f.t_start, f.path, f.size_bytes, d.verdict "
            f"FROM frame f LEFT JOIN detection d ON d.frame_id = f.id AND d.run_id = ? "
            f"WHERE f.t_start >= ? AND f.t_start < ? AND f.satellite IN ({marks})",
            (run_id, t0, t1, *satellites),
        )
    ]


def archive_map(rows, root, wavelengths):
    """slot -> absolute path, leaving out stub downloads (they hold no image)."""
    out = {}
    for satellite, wavelength, t_start, relpath, size, *_ in rows:
        if wavelength not in wavelengths:
            continue
        if size is not None and size < dataset.MIN_FRAME_BYTES:
            continue
        out[(satellite, wavelength, t_start)] = paths.abspath(relpath, root)
    return out


def judged_bad(rows, archive):
    """(bad slots, unjudged slots) among the frames in `archive`."""
    bad, unjudged = set(), set()
    for satellite, wavelength, t_start, _, _, verdict in rows:
        slot = (satellite, wavelength, t_start)
        if slot not in archive:
            continue
        if verdict is None:
            unjudged.add(slot)
        elif verdict == detectors.BAD:
            bad.add(slot)
    return bad, unjudged


class FlowFiller:
    """Per-band optical flow between the band's nearest clean frames either side.

    The classical repair the learned model is compared against, using the bench's
    own `new` variant so the two cannot differ in how they choose the anchors.
    """

    def __init__(self, overlay, wavelengths, bad):
        self.overlay = overlay
        self.wavelengths = tuple(wavelengths)
        self.bad = bad
        self.cache = {}

    def fill(self, satellite, when, bands):
        """{band index: array} for `bands`, or None if any of them has no anchor."""
        if len(self.cache) > FLOW_CACHE:
            self.cache.clear()
        out = {}
        for index in bands:
            slot = (satellite, self.wavelengths[index], when)
            image, ok = bench._replacement_band("new", self.overlay, slot, self.bad,
                                                self.cache, self.wavelengths)
            if not ok:
                return None
            out[index] = image
        return out


def make_repair(filler, overlay, satellites, wavelengths, bad):
    """``repair(satellite, when, bands)`` -> indexable by band, or None."""
    if filler == "optical_flow":
        return FlowFiller(overlay, wavelengths, bad).fill
    # The model reconstructs all six bands jointly, whichever were asked for; a band
    # no frame in its stack holds comes back NaN, and the slot is unfillable.
    joint = bench._LearnedFiller(overlay, satellites, wavelengths, bad)

    def repair(satellite, when, bands):
        filled = joint.fill(satellite, when)
        if filled is None or any(np.isnan(filled[band]).all() for band in bands):
            return None
        return filled

    return repair


def bands_to_fill(overlay, bad, satellite, when, wavelengths):
    """Indices of the bands at (satellite, when) that the model must supply."""
    return [index for index, wavelength in enumerate(wavelengths)
            if (satellite, wavelength, when) in bad
            or overlay.path((satellite, wavelength, when)) is None]


def render_day(day, satellite, satellites, wavelengths, rows, root, state, out,
               workers, batch, max_dual_gap, filler="learned"):
    """Composite every slot of one day for one satellite.  Returns a summary."""
    t0, t1 = window_for(day)
    archive = archive_map(rows, root, wavelengths)
    bad, unjudged = judged_bad(rows, archive)
    if unjudged:
        raise SystemExit(f"{day_name(day)}: {len(unjudged)} frames in the context "
                         f"window have no stored verdict; run `filter_FITS.py "
                         f"judge` over {day_name(t0)}..{day_name(t1)} first")

    overlay = cases.Overlay(archive=archive)
    unfillable = bench.unfillable_by_band(overlay, satellites, wavelengths, bad,
                                          cases.timeline(t0, t1), max_dual_gap)
    repair = make_repair(filler, overlay, satellites, wavelengths, bad)

    out_dir = os.path.join(out, filler, f"goes{satellite}")
    os.makedirs(out_dir, exist_ok=True)
    summary = {"written": 0, "filled": 0, "policy_gap": 0, "unfillable": 0,
               "existing": 0, "bad_bands": 0}
    jobs = []
    for when in cases.timeline(day, day + DAY):
        target = os.path.join(out_dir, f"Composite-{int(when)}.jpg")
        if os.path.exists(target):
            summary["existing"] += 1          # composites are immutable
            continue
        if when in unfillable:
            summary["policy_gap"] += 1        # the gap stays a gap: black frame
            continue
        missing = bands_to_fill(overlay, bad, satellite, when, wavelengths)
        filled = repair(satellite, when, missing) if missing else {}
        if filled is None:
            summary["unfillable"] += 1        # no usable source for a missing band
            continue
        arrays = []
        for index, wavelength in enumerate(wavelengths):
            if index in missing:
                arrays.append(filled[index])
                continue
            image, _ = fitsio.read_image(overlay.path((satellite, wavelength, when)))
            if image is None:
                # Judged good but unreadable now: the filler supplies it instead.
                extra = repair(satellite, when, [index])
                if extra is None:
                    arrays = None
                    break
                image = extra[index]
            arrays.append(image)
        if arrays is None:
            summary["unfillable"] += 1
            continue
        if missing:
            summary["filled"] += 1
            summary["bad_bands"] += len(missing)
        jobs.append((arrays, when, target))
        if len(jobs) >= batch:
            bench._flush_renders(jobs, workers)
            summary["written"] += len(jobs)
            jobs.clear()
    if jobs:
        bench._flush_renders(jobs, workers)
        summary["written"] += len(jobs)
    return summary


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--start", type=bench.parse_date, required=True)
    parser.add_argument("--end", type=bench.parse_date, required=True,
                        help="exclusive")
    parser.add_argument("--satellites", default="16,18",
                        help="the pair the model fuses; both are used as context")
    parser.add_argument("--render", default=None,
                        help="satellites to write composites for, in order "
                             "(default: all of --satellites)")
    parser.add_argument("--out", required=True)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--batch", type=int, default=12,
                        help="composites per worker dispatch; bounds peak memory")
    parser.add_argument("--max-dual-gap", type=int, default=availability.MAX_DUAL_GAP)
    parser.add_argument("--filler", choices=FILLERS, default="learned",
                        help="what repairs bad bands: the model (SUVI_MODEL) or "
                             "per-band optical flow")
    parser.add_argument("--run", type=int, default=None,
                        help="judgment run id (default: the run made under the current rules)")
    parser.add_argument("--db", default=None)
    args = parser.parse_args(argv)

    if args.filler == "learned" and fillers.load_learned() is None:
        raise SystemExit(f"needs a checkpoint: set {fillers.LEARNED_CHECKPOINT_ENV}")
    fitsio.quiet_astropy()
    satellites = tuple(bench.parse_int_list(args.satellites))
    render = tuple(bench.parse_int_list(args.render)) if args.render else satellites
    wavelengths = tuple(paths.WAVELENGTHS)
    conn = db.connect(args.db, readonly=True)
    run_id = args.run or detectors.judgment_run(conn)
    if run_id is None:
        raise SystemExit("No production judgment under the current rules; run "
                         "`filter_FITS.py judge` over the dates (and context) first.")
    root = paths.data_root()
    state = State(args.out, args.filler)
    days = day_starts(args.start, args.end)
    print(f"real run: {day_name(days[0])} .. {day_name(days[-1])}, "
          f"{len(days)} days, rendering {render} with {args.filler}, "
          f"context {satellites}, judgment run {run_id}", flush=True)

    started = time.time()
    for satellite in render:
        for index, day in enumerate(days, 1):
            if state.is_done(day, satellite):
                continue
            t0, t1 = window_for(day)
            rows = frame_rows(conn, t0, t1, satellites, run_id)
            began = time.time()
            summary = render_day(day, satellite, satellites, wavelengths, rows, root,
                                 state, args.out, args.workers, args.batch,
                                 args.max_dual_gap, args.filler)
            state.mark_done(day, satellite, summary)
            print(f"  g{satellite} [{index:>3}/{len(days)}] {day_name(day)}: "
                  f"{summary['written']} written ({summary['filled']} filled, "
                  f"{summary['bad_bands']} bands), {summary['policy_gap']} left "
                  f"black by policy, {summary['unfillable']} unfillable, "
                  f"{time.time() - began:.0f}s", flush=True)
    print(f"done in {(time.time() - started) / 3600:.1f} h")
    return 0


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