#!/usr/bin/env python
"""Scan rendered composites for defects that reached the output.

    render_qa.py renders/compare_v8/new-detection/learned/goes16 \\
        renders/compare_v8/new-detection/optical_flow/goes16 \\
        --periods 2024-10-03..2024-10-06 --out renders/compare_v8/qa

Detection is scored on labelled archive frames (`detector_eval.py`), but what the
viewer sees is the composite, and a defect can reach it without any frame-level
check failing: a judged-good frame that is bad, or a repair built from bad
neighbours (2024-01-01 16:28 in the rejected render: a missing 131 A band filled
from two static frames came out as static).  This reads the JPEGs themselves.

Each composite holds six band panels and the composite in the middle
(`merger_FITS.composite_from_arrays`).  For each of the seven, on luminance:

* **far_field** -- 99th percentile beyond 1.6 solar radii.  The colour maps clip
  the quiet far field to black, so bars, blocks, ghost discs and static light it.
* **noise** -- spread of the high-pass residual there: static's signature.
* **spike** -- how far the panel stands out of line with the slots either side,
  per pixel ``max(0, |x - (prev+next)/2| - |next-prev|/2)``.  Zero for a ramp or
  a step, so flare onsets and rotation pass; a one-frame defect or a bad repair
  does not.  Needs both neighbours one cadence away.

A frame is flagged when any panel exceeds its limit (`LIMITS`).  The flagged
frames are listed per stream, worst first, with contact sheets for review.
"""

import argparse
import datetime
import json
import os
import sys
from multiprocessing import Pool

import numpy as np
from PIL import Image, ImageDraw

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

import merger_FITS
from make_comparison_video import build_timeline, parse_periods, stream_frames
from suvi import paths

#: Panel order: the six bands as merger_FITS lays them out, then the composite.
PANELS = tuple(name.lstrip("0").rstrip("A") for name in merger_FITS.image_names) + ("composite",)
#: Trimmed band size, and the shrunken copies along both sides.
FRAME_W = 1280 - 2 * merger_FITS.TRIM_X
FRAME_H = 1280 - 2 * merger_FITS.TRIM_Y
PANEL_W, PANEL_H = FRAME_W // 3, FRAME_H // 3
CANVAS = (FRAME_W + 2 * PANEL_W, FRAME_H)
#: Largest solar radius in a 1280-pixel frame (perihelion, 2.5"/px), and where
#: the far field starts, as in `detectors.field_v1`.
SUN_RADIUS = 392
FAR_FIELD_RADIUS = 1.6
#: Rows of the composite under the title; its timestamp changes every frame.
TITLE_ROWS = 60
#: Downsampling for the spike test: JPEG blocks and grain average out.
SPIKE_BIN = 4

#: Per-panel limits on luminance (0..1).  Calibrated on the rejected v7_s128
#: GOES-16 render over 2024-01-01..07, 05-08..10 and 10-03..05 (4,680
#: composites) against the 46 slots holding frames the stored judgment rejects
#: but that render's detection passed.  38 were flagged -- static at noise
#: 0.02-0.06, bars and blocks at far field 0.2-1.0, the rest at spike 0.02 or
#: more; the other 8 are 94 A frames the series check rejects that look normal in
#: the panel.  The one other flag was a real bad repair (16:28 above).  Unflagged
#: frames, X-class flares included, peaked at far field 0.004, noise 0.002 and
#: spike 0.006.
LIMITS = {"far_field": 0.15, "noise": 0.005, "spike": 0.015}
METRICS = tuple(LIMITS)


def panel_boxes():
    """(x0, y0, x1, y1) of each panel in `PANELS` order."""
    boxes = []
    for index in range(6):
        x0 = (index % 2) * (CANVAS[0] - PANEL_W)
        y0 = (index // 2) * PANEL_H
        boxes.append((x0, y0, x0 + PANEL_W, y0 + PANEL_H))
    boxes.append((PANEL_W, 0, PANEL_W + FRAME_W, FRAME_H))
    return tuple(boxes)


def far_field_masks():
    """Boolean far-field mask per panel, in `PANELS` order."""
    masks = []
    for x0, y0, x1, y1 in panel_boxes():
        height, width = y1 - y0, x1 - x0
        scale = width / FRAME_W
        yy, xx = np.mgrid[0:height, 0:width]
        radius = np.hypot(xx - (width - 1) / 2, yy - (height - 1) / 2)
        mask = radius > FAR_FIELD_RADIUS * SUN_RADIUS * scale
        if (x0, y0) == (PANEL_W, 0):
            mask[:TITLE_ROWS] = False
        masks.append(mask)
    return tuple(masks)


_MASKS = []


def _box_blur(image):
    padded = np.pad(image, 1, mode="edge")
    height, width = image.shape
    return sum(padded[i:i + height, j:j + width]
               for i in range(3) for j in range(3)) / 9.0


def _bin(image):
    height = image.shape[0] // SPIKE_BIN * SPIKE_BIN
    width = image.shape[1] // SPIKE_BIN * SPIKE_BIN
    return image[:height, :width].reshape(
        height // SPIKE_BIN, SPIKE_BIN, width // SPIKE_BIN, SPIKE_BIN).mean(axis=(1, 3))


def measure(path):
    """(far_field, noise) per panel and the binned panels for the spike test, or
    None for an unreadable or wrongly sized composite."""
    if not _MASKS:
        _MASKS.extend(far_field_masks())
    try:
        with Image.open(path) as image:
            if image.size != CANVAS:
                return None
            luminance = np.asarray(image.convert("L"), dtype=np.float32) / 255.0
    except OSError:
        return None
    levels, binned = [], []
    for (x0, y0, x1, y1), mask in zip(panel_boxes(), _MASKS):
        panel = luminance[y0:y1, x0:x1]
        if (x0, y0) == (PANEL_W, 0):
            panel = panel.copy()
            panel[:TITLE_ROWS] = 0.0          # the title is the same blank in all three
        far = panel[mask]
        residual = (panel - _box_blur(panel))[mask]
        levels.append((float(np.percentile(far, 99)), float(residual.std())))
        binned.append(_bin(panel))
    return levels, binned


def spike(previous, current, following):
    """Mean out-of-line excess of `current` against its two neighbours."""
    midpoint = (previous + following) / 2.0
    excess = np.abs(current - midpoint) - np.abs(following - previous) / 2.0
    return float(np.maximum(excess, 0.0).mean())


def scan(frames, timeline, workers=2):
    """{time: {"far_field": [...7], "noise": [...7], "spike": [...7] or None}}
    for every slot of `timeline` with a readable composite."""
    wanted = [when for when in timeline if when in frames]
    results = {}
    window = []                                   # last three (time, binned)
    with Pool(workers) as pool:
        measured = pool.imap(measure, [frames[when] for when in wanted], 8)
        for when, outcome in zip(wanted, measured):
            if outcome is None:
                window.clear()
                continue
            levels, binned = outcome
            results[when] = {"far_field": [level[0] for level in levels],
                             "noise": [level[1] for level in levels],
                             "spike": None}
            window = (window + [(when, binned)])[-3:]
            if len(window) == 3:
                (t0, a), (t1, b), (t2, c) = window
                if t1 - t0 == paths.CADENCE and t2 - t1 == paths.CADENCE:
                    results[t1]["spike"] = [spike(x, y, z) for x, y, z in zip(a, b, c)]
    return results


def flag(results, limits=LIMITS):
    """Flagged frames, worst first: (time, severity, [(panel, metric, value)])."""
    flagged = []
    for when, values in results.items():
        hits = []
        for metric, limit in limits.items():
            for panel, value in zip(PANELS, values[metric] or ()):
                if value > limit:
                    hits.append((panel, metric, value))
        if hits:
            severity = max(value / limits[metric] for _, metric, value in hits)
            flagged.append((when, severity, sorted(hits, key=lambda h: -h[2] / limits[h[1]])))
    flagged.sort(key=lambda item: (-item[1], item[0]))
    return flagged


def stream_name(directory):
    """'.../compare_v8/new-detection/learned/goes16' -> 'new-detection_learned_goes16'."""
    parts = os.path.normpath(os.path.abspath(directory)).split(os.sep)
    return "_".join(parts[-3:])


def _stamp(when):
    return datetime.datetime.fromtimestamp(when, datetime.timezone.utc).strftime("%Y-%m-%d %H:%M")


def contact_sheets(flagged, frames, out_prefix, columns=4, rows=6, limit=48):
    """Thumbnails of the worst flagged frames, offending panels outlined.
    Returns the paths written."""
    boxes = dict(zip(PANELS, panel_boxes()))
    scale = 4
    tile_w, tile_h = CANVAS[0] // scale, CANVAS[1] // scale
    chosen = flagged[:limit]
    per_sheet = columns * rows
    written = []
    for start in range(0, len(chosen), per_sheet):
        members = chosen[start:start + per_sheet]
        sheet_rows = (len(members) + columns - 1) // columns
        sheet = Image.new("RGB", (tile_w * columns, tile_h * sheet_rows))
        for index, (when, severity, hits) in enumerate(members):
            with Image.open(frames[when]) as image:
                tile = image.convert("RGB").resize((tile_w, tile_h))
            draw = ImageDraw.Draw(tile)
            for panel in {hit[0] for hit in hits}:
                x0, y0, x1, y1 = (v // scale for v in boxes[panel])
                draw.rectangle((x0, y0, x1 - 1, y1 - 1), outline=(255, 0, 0), width=2)
            panel, metric, value = hits[0]
            draw.rectangle((0, tile_h - 14, tile_w, tile_h), fill=(0, 0, 0))
            draw.text((3, tile_h - 13), f"{_stamp(when)}  {panel} {metric} {value:.3f}",
                      fill=(255, 255, 0))
            sheet.paste(tile, (tile_w * (index % columns), tile_h * (index // columns)))
        path = f"{out_prefix}_{start // per_sheet + 1:02d}.jpg"
        sheet.save(path, quality=85)
        written.append(path)
    return written


def report(name, results, flagged, limits=LIMITS):
    """The JSON record of one stream's scan."""
    return {
        "stream": name,
        "limits": limits,
        "frames": len(results),
        "flagged": [{"time": when, "utc": _stamp(when), "severity": round(severity, 2),
                     "hits": [{"panel": p, "metric": m, "value": round(v, 4)}
                              for p, m, v in hits]}
                    for when, severity, hits in flagged],
    }


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("streams", nargs="+", help="directories of Composite-*.jpg")
    parser.add_argument("--out", required=True, help="report and contact-sheet directory")
    parser.add_argument("--periods", default=None,
                        help="comma-separated UTC date ranges START..END (end exclusive)")
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--sheet", type=int, default=48,
                        help="worst flagged frames per stream to put on contact sheets")
    args = parser.parse_args(argv)

    periods = parse_periods(args.periods) if args.periods else None
    names = [stream_name(directory) for directory in args.streams]
    if len(set(names)) != len(names):
        raise SystemExit(f"streams must differ in their last three path parts: {names}")
    os.makedirs(args.out, exist_ok=True)
    summary = []
    for directory, name in zip(args.streams, names):
        frames = stream_frames(directory)
        if not frames:
            raise SystemExit(f"{directory}: no composites")
        results = scan(frames, build_timeline({name: frames}, periods), args.workers)
        flagged = flag(results)
        with open(os.path.join(args.out, f"{name}.json"), "w") as handle:
            json.dump(report(name, results, flagged), handle, indent=1)
        sheets = contact_sheets(flagged, frames, os.path.join(args.out, name),
                                limit=args.sheet)
        by_panel = {}
        for _, _, hits in flagged:
            for panel in {hit[0] for hit in hits}:
                by_panel[panel] = by_panel.get(panel, 0) + 1
        summary.append((name, len(results), len(flagged), by_panel, len(sheets)))
        print(f"{name}: {len(results)} composites, {len(flagged)} flagged", flush=True)
        for when, severity, hits in flagged[:10]:
            panel, metric, value = hits[0]
            print(f"    {_stamp(when)}  x{severity:.1f}  {panel} {metric} {value:.3f}"
                  + (f"  (+{len(hits) - 1} more)" if len(hits) > 1 else ""))

    print(f"\n{'stream':<40} {'composites':>10} {'flagged':>8}  by panel")
    for name, count, flagged, by_panel, sheets in summary:
        panels = ", ".join(f"{p} {by_panel[p]}" for p in PANELS if p in by_panel)
        print(f"{name:<40} {count:>10} {flagged:>8}  {panels}")
    print(f"\nreports and contact sheets in {args.out}")
    return 0


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