#!/usr/bin/env python3
"""
stretch_thermal.py - Contrast-stretch a session's radiometric thermal frames
to one common signal window, for consistent, comparable visualisation.

The thermal PNGs are 16-bit radiometric (raw signal counts, monotonic with
temperature but not calibrated to degC). Each frame on its own spans only a
narrow part of the 16-bit range and uses a different sub-range, so viewing the
raw files is near-black and frame-to-frame brightness flickers.

This computes ONE window [lo, hi] pooled across every frame in the session
(default: 1st-99th percentile, robust to outlier pixels), then maps that window
to 8-bit and writes the result with a chosen palette. Because all frames share
the same window, brightness/colour is directly comparable across the session.

Usage:
    python stretch_thermal.py [SESSION_DIR] [options]

    --colormap {inferno,ironbow,gray}   palette (default: inferno)
    --out DIR            output folder (default: <session>/thermal_stretched)
    --lo-pct P           low percentile  (default: 1.0)
    --hi-pct P           high percentile (default: 99.0)
    --lo N / --hi N      hard-set the signal window, overriding percentiles
    --absolute           use the true session min/max instead of percentiles
    --sample             write one side-by-side comparison of all palettes and exit
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import embed_metadata as embed  # noqa: E402  (reuse exiftool + tag building)
import session_map as smap       # noqa: E402  (session discovery)

import numpy as np            # noqa: E402
from PIL import Image         # noqa: E402

U16_MAX = 65536


# --------------------------------------------------------------------------- #
# palettes - built-in 256-entry LUTs, no matplotlib dependency
# --------------------------------------------------------------------------- #
def _lut_from_anchors(anchors) -> np.ndarray:
    """Build a 256x3 uint8 LUT by linear interpolation between (pos, r,g,b) anchors."""
    pos = np.array([a[0] for a in anchors])
    rgb = np.array([a[1:] for a in anchors], dtype=float)
    x = np.linspace(0.0, 1.0, 256)
    lut = np.stack([np.interp(x, pos, rgb[:, c]) for c in range(3)], axis=1)
    return np.clip(lut, 0, 255).astype(np.uint8)


INFERNO = _lut_from_anchors([
    (0.00, 0, 0, 4), (0.13, 31, 12, 72), (0.25, 85, 15, 109), (0.38, 136, 34, 106),
    (0.50, 186, 54, 85), (0.63, 227, 89, 51), (0.75, 249, 140, 10),
    (0.88, 249, 201, 50), (1.00, 252, 255, 164),
])
IRONBOW = _lut_from_anchors([
    (0.00, 0, 0, 0), (0.12, 0, 0, 70), (0.25, 60, 0, 130), (0.40, 160, 0, 120),
    (0.55, 220, 40, 60), (0.70, 250, 110, 0), (0.85, 255, 200, 40), (1.00, 255, 255, 255),
])
PALETTES = {"inferno": INFERNO, "ironbow": IRONBOW}


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def find_session(arg: str | None) -> Path:
    if arg:
        return Path(arg).resolve()
    sessions = smap.list_sessions(Path(__file__).resolve().parent)
    if not sessions:
        sys.exit("No session folder found nearby; pass SESSION_DIR explicitly.")
    return max(sessions, key=lambda p: p.stat().st_mtime)


def list_frames(session: Path) -> list[Path]:
    """Sorted thermal PNG frames in a session (raises if none)."""
    thermal_dir = Path(session) / "thermal"
    if not thermal_dir.is_dir():
        raise RuntimeError(f"No thermal/ folder in {session}")
    pngs = sorted(thermal_dir.glob("*.png"))
    if not pngs:
        raise RuntimeError(f"No PNG frames in {thermal_dir}")
    return pngs


def session_histogram(pngs, progress=None) -> np.ndarray:
    """Pooled value histogram over all frames (exact for integer signal data)."""
    hist = np.zeros(U16_MAX, dtype=np.int64)
    n = len(pngs)
    for i, p in enumerate(pngs, 1):
        arr = np.asarray(Image.open(p)).astype(np.uint16, copy=False)
        hist += np.bincount(arr.ravel(), minlength=U16_MAX)
        if progress:
            progress(0.5 * i / n, f"Analysing frame {i}/{n}")
    return hist


def percentile_from_hist(hist: np.ndarray, pct: float) -> int:
    total = hist.sum()
    cdf = np.cumsum(hist)
    return int(np.searchsorted(cdf, total * pct / 100.0))


def compute_window(pngs, lo_pct=1.0, hi_pct=99.0, lo=None, hi=None,
                   absolute=False, progress=None) -> dict:
    """Resolve the common [lo, hi] signal window for a set of frames."""
    hist = session_histogram(pngs, progress=progress)
    nz = np.nonzero(hist)[0]
    abs_lo, abs_hi = int(nz[0]), int(nz[-1])
    if lo is not None or hi is not None:
        rlo = lo if lo is not None else abs_lo
        rhi = hi if hi is not None else abs_hi
        how = "manual"
    elif absolute:
        rlo, rhi, how = abs_lo, abs_hi, "absolute min/max"
    else:
        rlo = percentile_from_hist(hist, lo_pct)
        rhi = percentile_from_hist(hist, hi_pct)
        how = f"{lo_pct:g}-{hi_pct:g} percentile"
    if rhi <= rlo:
        rhi = rlo + 1
    return {"lo": rlo, "hi": rhi, "abs_lo": abs_lo, "abs_hi": abs_hi, "how": how}


def render(arr: np.ndarray, lo: int, hi: int, colormap: str) -> Image.Image:
    """Stretch one 16-bit frame through [lo, hi] -> 8-bit, apply palette."""
    span = max(hi - lo, 1)
    norm = np.clip((arr.astype(np.float32) - lo) / span, 0.0, 1.0)
    idx = (norm * 255.0 + 0.5).astype(np.uint8)
    if colormap == "gray":
        return Image.fromarray(idx, mode="L")
    return Image.fromarray(PALETTES[colormap][idx], mode="RGB")


def render_frame(path, lo: int, hi: int, colormap: str) -> Image.Image:
    """Open one PNG and render it with the given window/palette (for previews)."""
    return render(np.asarray(Image.open(path)).astype(np.uint16), lo, hi, colormap)


def make_sample_image(pngs, lo: int, hi: int, gap: int = 8) -> Image.Image:
    """Side-by-side comparison of all palettes for the middle frame."""
    mid = np.asarray(Image.open(pngs[len(pngs) // 2])).astype(np.uint16)
    h, w = mid.shape
    tiles = [render(mid, lo, hi, c).convert("RGB") for c in ("gray", "inferno", "ironbow")]
    canvas = Image.new("RGB", (w * 3 + gap * 2, h), (255, 255, 255))
    for i, t in enumerate(tiles):
        canvas.paste(t, (i * (w + gap), 0))
    return canvas


def _embed_exif(pngs, out_dir: Path, exiftool=None, log=None):
    """Write each frame's JSON sidecar into the matching stretched PNG (GPS, true
    camera orientation, time, full JSON) so the 8-bit outputs are self-contained."""
    log = log or (lambda *_: None)
    try:
        et = exiftool or embed.ensure_exiftool(None, Path(__file__).resolve().parent / "tools")
    except BaseException as e:  # noqa: BLE001  (ensure_exiftool may sys.exit)
        log(f"  ! skipping EXIF embed - exiftool unavailable ({e})")
        return
    arg_lines, dsts = [], []
    for p in pngs:
        sidecar = p.with_suffix(".json")
        if not sidecar.exists():
            continue
        try:
            data = json.loads(sidecar.read_text(encoding="utf-8"))
        except Exception:  # noqa: BLE001
            continue
        dst = out_dir / (p.stem + ".png")
        dsts.append(dst)
        arg_lines += ["-overwrite_original", *embed.build_tags(data, {}, None, "gps"),
                      str(dst), "-execute"]
    if not arg_lines:
        return
    argfile = out_dir / "_exif_args.txt"
    argfile.write_text("\n".join(arg_lines) + "\n", encoding="utf-8")
    cmd = [et, "-m", "-charset", "UTF8", "-charset", "filename=UTF8", "-@", str(argfile)]
    samples = {dsts[0], dsts[len(dsts) // 2], dsts[-1]}   # spot-check first/mid/last

    def gps_written():
        return all(subprocess.run([et, "-s", "-s", "-s", "-GPSLatitude", str(s)],
                                  capture_output=True, text=True).stdout.strip() for s in samples)

    log("Embedding GPS / orientation EXIF into the stretched frames ...")
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode == 0 and not gps_written():
        log("  ! GPS not present after first pass - retrying once ...")
        proc = subprocess.run(cmd, capture_output=True, text=True)
    argfile.unlink(missing_ok=True)

    err = (proc.stderr or "").strip()
    if err:
        log(err)
    if proc.returncode != 0 or not gps_written():
        raise RuntimeError(
            "Could not embed GPS/EXIF into the stretched PNGs (exiftool exit "
            f"{proc.returncode}). The usual cause is a FULL DISK - free up space and re-run. "
            "The stretched images were written; only the embedded metadata is missing."
            + (f"\nexiftool: {err}" if err else ""))
    log(f"EXIF embedded into {len(dsts)} stretched frames.")


def stretch_session(session, out_dir=None, colormap="inferno", lo=None, hi=None,
                    embed_exif=True, exiftool=None, progress=None, log=None) -> dict:
    """
    Stretch every thermal frame of a session through the [lo, hi] window and
    write 8-bit images with the chosen palette. lo/hi must already be resolved
    (see compute_window). If embed_exif, the JSON sidecar metadata (GPS + true
    camera orientation) is written into each output PNG. Callbacks optional.
    """
    session = Path(session)
    log = log or (lambda *_: None)
    progress = progress or (lambda *_: None)
    pngs = list_frames(session)
    out_dir = Path(out_dir) if out_dir else session / "thermal_stretched"
    out_dir.mkdir(parents=True, exist_ok=True)

    n = len(pngs)
    log(f"Writing {n} {colormap} frames to {out_dir} ...")
    for i, p in enumerate(pngs, 1):
        render_frame(p, lo, hi, colormap).save(out_dir / (p.stem + ".png"))
        progress(0.5 + 0.45 * i / n, f"Rendering {i}/{n}")
    if embed_exif:
        _embed_exif(pngs, out_dir, exiftool=exiftool, log=log)
    progress(1.0, "Done")
    log(f"Done. {n} frames written.\nOutput: {out_dir}")
    return {"out_dir": out_dir, "frames": n}


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #
def main() -> int:
    ap = argparse.ArgumentParser(description="Stretch session thermal frames to a common window.")
    ap.add_argument("session", nargs="?", help="session folder (default: newest session_* here)")
    ap.add_argument("--colormap", choices=["inferno", "ironbow", "gray"], default="inferno")
    ap.add_argument("--out", help="output folder (default: <session>/thermal_stretched)")
    ap.add_argument("--lo-pct", type=float, default=1.0, help="low percentile (default 1.0)")
    ap.add_argument("--hi-pct", type=float, default=99.0, help="high percentile (default 99.0)")
    ap.add_argument("--lo", type=int, help="hard low signal value (overrides percentile)")
    ap.add_argument("--hi", type=int, help="hard high signal value (overrides percentile)")
    ap.add_argument("--absolute", action="store_true", help="use true session min/max")
    ap.add_argument("--no-embed-exif", action="store_true",
                    help="do not write GPS/orientation EXIF into the stretched PNGs")
    ap.add_argument("--sample", action="store_true",
                    help="write one side-by-side palette comparison and exit")
    args = ap.parse_args()

    session = find_session(args.session)
    pngs = list_frames(session)
    win = compute_window(pngs, args.lo_pct, args.hi_pct, args.lo, args.hi, args.absolute)
    lo, hi = win["lo"], win["hi"]
    print(f"Session : {session}")
    print(f"Frames  : {len(pngs)}  | absolute signal range {win['abs_lo']}-{win['abs_hi']}")
    print(f"Window  : {lo}-{hi}  ({win['how']})")

    if args.sample:
        out = session / "thermal_palette_sample.png"
        make_sample_image(pngs, lo, hi).save(out)
        print(f"\nSample written (left->right: gray | inferno | ironbow):\n  {out}")
        return 0

    out_dir = Path(args.out).resolve() if args.out else None
    res = stretch_session(session, out_dir, args.colormap, lo, hi,
                          embed_exif=not args.no_embed_exif, log=print)
    print(f"\nWrote {res['frames']} {args.colormap} frames to:\n  {res['out_dir']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())