#!/home/robert/weather-env/bin/python
"""Persistent warmest-pixel SST composite (GOES-East CONUS, 2 km) — runs HOURLY.

Each run merges the SST of only the LAST HOUR's scans into a per-pixel running-
max grid that is PERSISTED to .npy between runs, then renders the composite.
The grid RESETS daily at 23:00 UTC (a fresh accumulation window begins), so the
image is the warmest clear-sky SST seen since the last 23Z — clouds (cold) get
filtered out as they drift over the day.  Sequential merge -> low RAM.

Reuses the tile-assembly / border / colorbar helpers from make_goes.py.
Schedule:  0 * * * *  /home/robert/weather-env/bin/python .../make_sst.py
"""
import os, time, json
from datetime import datetime, timedelta
import numpy as np
from PIL import Image
from matplotlib import colormaps, colors

import make_goes as G   # shared helpers (same directory)

C14, C15 = "TIR/TIRE14", "TIR/TIRE15"
STATE_NPY  = os.path.join(G.OUT_DIR, "sst_maxgrid.npy")
STATE_JSON = os.path.join(G.OUT_DIR, "sst_state.json")
OUTFILE    = "goes_east_conus_sst.png"
WINDOW_MIN = 60         # merge only scans from the last hour
RESET_HOUR = 23         # daily reset at 23:00 (box local = UTC)
VMIN, VMAX = 10.0, 32.0


def _sst(T11, T12):
    sst = 1.0346*T11 + 2.5779*(T11 - T12) - 283.21               # MCSST split-window, degC
    bad = ~np.isfinite(sst) | (T11 < 271.0) | (sst < -2) | (sst > 36)   # clouds/land/range
    return np.where(bad, -np.inf, sst).astype(np.float32)        # invalid never wins a max


def _window_start():
    """Epoch of the most recent RESET_HOUR boundary (box-local = UTC)."""
    now = datetime.now()
    b = now.replace(hour=RESET_HOUR, minute=0, second=0, microsecond=0)
    if now < b:
        b -= timedelta(days=1)
    return b.timestamp()


def main():
    ws = _window_start()
    g14 = G._group_ts(os.path.join(G.ROOT, C14))
    g15 = G._group_ts(os.path.join(G.ROOT, C15))
    stamps = sorted(set(g14) & set(g15))
    if not stamps:
        print("no SST scans on disk"); return
    ts_mt = {ts: max(os.path.getmtime(f) for f in g14[ts] + g15[ts]) for ts in stamps}
    newest = max(ts_mt.values())
    recent = [ts for ts in stamps if ts_mt[ts] >= newest - WINDOW_MIN*60]   # last hour only

    # load the persisted grid unless a new daily (23Z) window has started -> reset
    state = {}
    try:
        state = json.load(open(STATE_JSON))
    except Exception:
        pass
    mx = None
    if state.get("window_start") == ws and os.path.exists(STATE_NPY):
        try:
            mx = np.load(STATE_NPY)
        except Exception:
            mx = None
    reset = mx is None

    ref = None; merged = 0
    for ts in recent:
        a14 = G._assemble(g14[ts]); a15 = G._assemble(g15[ts])
        if not (a14 and a15):
            continue
        sst = _sst(a14["canvas"], G._resize(a15["canvas"], a14["shape"]))
        if mx is None or mx.shape != a14["shape"]:
            mx = np.full(a14["shape"], -np.inf, np.float32)
        np.maximum(mx, sst, out=mx)                              # running max in place
        ref = a14; merged += 1
        del a14, a15, sst
    if mx is None:
        print("no scans assembled"); return

    np.save(STATE_NPY, mx)
    json.dump({"window_start": ws, "updated": time.time(), "shape": list(mx.shape)},
              open(STATE_JSON, "w"))

    if ref is None:                                              # no new scans -> re-render persisted grid
        ref = G.assemble(C14)
        if ref is None:
            print("no grid metadata"); return

    masked = ~np.isfinite(mx)
    cmap = colormaps["turbo"]
    rgb = (cmap(colors.Normalize(VMIN, VMAX)(np.where(masked, np.nan, mx)))[..., :3] * 255).astype(np.uint8)
    rgb[masked] = 0
    img = Image.fromarray(rgb, "RGB")
    d = G._overlay(img, ref, (255, 255, 255), land_black=True)
    G._colorbar(img, cmap, VMIN, VMAX, "SST degC")
    since = time.strftime("%m-%d %HZ", time.gmtime(ws))
    G._label(d, ref, f"GOES-East SST warmest-pixel composite since {since}  "
                     f"(CONUS 2km, +{merged} scans{', RESET' if reset else ''})", (255, 255, 255))
    G._save(img, OUTFILE)
    print(f"window_start={since} merged={merged} reset={reset}")


if __name__ == "__main__":
    main()