#!/usr/bin/env python3
"""Renders the example scenes at two revisions and compares them, for
scripts/serve-gallery.py's Delta tab: what a change did to the images.

    python3 scripts/gallery-delta.py [--base REV] [--head REV] [--noise] [--scenes a,b]

With no revisions: with uncommitted changes, HEAD against the working tree; with none,
HEAD~1 against HEAD. A revision is a commit (anything git rev-parse takes) or
`worktree`, the working tree as it is.

Each revision renders every example scene in its own examples/scenes/ at the gallery's
tiers (lighting L0 and L1, camera C1), the pretty preview for looking at and the raw DN
for comparing. Renders are noiseless (`resolved render --noiseless`), so an unchanged
scene comes out identical and a changed one differs only where its light does; --noise
keeps the detector's noise, which then also differs wherever the light does. A commit
is built in its own git worktree under target/delta/, with data/ and the fetched models
linked in, and its renders are kept there by commit, so comparing the next commit
reuses them. The working tree is built and rendered fresh each time.

The comparison goes to gallery/delta/: before, after and difference images, and
summary.json with the change in DN per image. The difference image is the after image
dimmed to gray, with every pixel that changed colored by how much on a log scale:
orange where it got brighter, blue where it got darker.

Needs numpy, Pillow and ImageMagick (`magick`, to read 16-bit PNGs, which Pillow reads
as 8-bit). The first build of a commit takes a few minutes; later ones are incremental.
After deleting target/, `git worktree prune` forgets the worktrees.
"""

import argparse
import concurrent.futures
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
DELTA = ROOT / "target" / "delta"
OUT = ROOT / "gallery" / "delta"
TIERS = ["L0", "L1"]
CAMERA = "C1"
# Renders at once. Each loads its own copy of the Earth and Moon maps; four at once ran
# the RTX 3060 (12 GB) out of memory, and with a sharp Earth map (4.9 GB) two do.
WORKERS = 1 if any((ROOT / "data" / "earth").glob("sharp-*.tiles")) else 2
BRIGHTER = np.array([255, 150, 40], dtype=float)
DARKER = np.array([60, 140, 255], dtype=float)


def git(*args, cwd=ROOT):
    return subprocess.run(
        ["git", *args], cwd=cwd, check=True, capture_output=True, text=True
    ).stdout.strip()


def dirty():
    return bool(git("status", "--porcelain"))


class Side:
    """One revision: where its sources are, how to build it, where its renders go."""

    def __init__(self, rev, slot):
        if rev == "worktree":
            self.sha = None
            self.label = "working tree"
            self.root = ROOT
            self.target = ROOT / "target"
            self.renders = DELTA / "renders" / "worktree"
        else:
            self.sha = git("rev-parse", "--verify", f"{rev}^{{commit}}")
            self.label = git("log", "-1", "--format=%h %s", self.sha)
            self.root = DELTA / slot
            self.target = DELTA / f"{slot}-target"
            self.renders = DELTA / "renders" / self.sha

    def prepare(self):
        """Checks the revision out into its worktree, with the untracked data linked."""
        if self.sha is None:
            return
        git("worktree", "prune")
        if (self.root / ".git").exists():
            git("checkout", "--detach", "--force", self.sha, cwd=self.root)
        else:
            DELTA.mkdir(parents=True, exist_ok=True)
            git("worktree", "add", "--detach", "--force", str(self.root), self.sha)
        links = {
            self.root / "data": ROOT / "data",
            self.root / "examples" / "models" / "nasa": ROOT / "examples" / "models" / "nasa",
            # The library's astrokit path dependency, ../../../astrokit from its crate.
            DELTA / "astrokit": (ROOT / ".." / "astrokit").resolve(),
        }
        for link, target in links.items():
            if target.exists() and not link.exists():
                link.symlink_to(target)

    def build(self):
        env = dict(os.environ, CARGO_TARGET_DIR=str(self.target))
        print(f"building {self.label}", flush=True)
        subprocess.run(
            ["cargo", "build", "--quiet", "--release", "-p", "resolved-cli"],
            cwd=self.root,
            env=env,
            check=True,
        )
        self.binary = self.target / "release" / "resolved"
        usage = subprocess.run(
            [self.binary, "render", "--help"], capture_output=True, text=True
        ).stdout
        self.has_noiseless = "--noiseless" in usage

    def scenes(self):
        return {p.stem: p for p in sorted((self.root / "examples" / "scenes").glob("*.toml"))}

    def render_all(self, names, noise, fresh):
        """Renders each scene at each tier into its own directory, skipping those
        already rendered at this commit (failures are tried again). Returns
        {(scene, tier): directory or error}."""
        mode = "noise" if noise else "noiseless"
        base = self.renders / mode
        if self.sha is None or fresh:
            shutil.rmtree(base, ignore_errors=True)
        scenes = self.scenes()
        jobs = [(n, t) for n in names if n in scenes for t in TIERS]

        def one(job):
            name, tier = job
            out = base / f"{name}-{tier}"
            if (out / "dn.png").is_file():
                return job, out
            out.mkdir(parents=True, exist_ok=True)
            args = [
                self.binary, "render", "--scene", str(scenes[name].relative_to(self.root)),
                "--lighting", tier, "--camera", CAMERA, "--pretty",
                "--preview", str(out / "pretty.png"), "--out", str(out / "dn.png"),
            ]
            if not noise:
                args.append("--noiseless")
            result = subprocess.run(args, cwd=self.root, capture_output=True, text=True)
            if result.returncode != 0:
                return job, result.stderr.strip().removeprefix("error: ")
            return job, out

        print(f"rendering {self.label}: {len(jobs)} images", flush=True)
        with concurrent.futures.ThreadPoolExecutor(WORKERS) as pool:
            return dict(pool.map(one, jobs))


def read_dn(path):
    """A 16-bit PNG as DN: its values shifted down past the zero padding below them."""
    image = Image.open(path)
    width, height = image.size
    channels = len(image.getbands())
    kind = "rgb" if channels == 3 else "gray"
    raw = subprocess.run(
        ["magick", str(path), "-depth", "16", "-endian", "LSB", f"{kind}:-"],
        check=True,
        capture_output=True,
    ).stdout
    values = np.frombuffer(raw, "<u2").reshape(height, width, -1)
    return values


def padding(values):
    """The number of zero bits below every value: 4 for 12-bit DN in 16 bits."""
    bits = int(np.bitwise_or.reduce(values, axis=None))
    return (bits & -bits).bit_length() - 1 if bits else 0


def difference_image(after_pretty, delta):
    """The after image dimmed to gray, with changed pixels colored on a log scale."""
    gray = np.asarray(Image.open(after_pretty).convert("L"), dtype=float)[..., None]
    picture = np.repeat(gray * 0.3, 3, axis=2)
    size = np.abs(delta).max(axis=2).astype(float)
    largest = size.max()
    if largest > 0:
        changed = size > 0
        strength = 0.35 + 0.65 * np.log1p(size) / np.log1p(largest)
        color = np.where((delta.sum(axis=2) >= 0)[..., None], BRIGHTER, DARKER)
        mix = (strength * changed)[..., None]
        picture = picture * (1 - mix) + color * mix
    return Image.fromarray(np.clip(picture, 0, 255).astype(np.uint8))


def compare(before, after, scene_changed, name, tier):
    item = {"scene": name, "tier": tier, "scene_changed": scene_changed}
    stem = f"{name}-{tier}"
    if isinstance(before, Path):
        shutil.copy(before / "pretty.png", OUT / f"{stem}-before.png")
        item["before"] = f"{stem}-before.png"
    if isinstance(after, Path):
        shutil.copy(after / "pretty.png", OUT / f"{stem}-after.png")
        item["after"] = f"{stem}-after.png"
    if not isinstance(before, Path) or not isinstance(after, Path):
        item["status"] = "new" if before is None else "removed" if after is None else "failed"
        for side, value in (("before", before), ("after", after)):
            if isinstance(value, str):
                item[f"{side}_error"] = value
        return item
    a, b = read_dn(before / "dn.png"), read_dn(after / "dn.png")
    if a.shape != b.shape:
        item.update(status="resized", shape=[list(a.shape), list(b.shape)])
        return item
    shift = min(padding(a), padding(b))
    delta = (b.astype(np.int32) >> shift) - (a.astype(np.int32) >> shift)
    size = np.abs(delta)
    changed = (size.max(axis=2) > 0).mean()
    item.update(
        status="changed" if changed > 0 else "same",
        pixels=int((size.max(axis=2) > 0).sum()),
        max_dn=int(size.max()),
        mean_dn=float(size.mean()),
        changed=float(changed),
        brighter=float((delta.sum(axis=2) > 0).mean()),
        darker=float((delta.sum(axis=2) < 0).mean()),
    )
    if changed > 0:
        difference_image(after / "pretty.png", delta).save(OUT / f"{stem}-diff.png")
        item["diff"] = f"{stem}-diff.png"
    return item


def main():
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--base", help="the revision before (default: see above)")
    parser.add_argument("--head", help="the revision after, or worktree")
    parser.add_argument("--noise", action="store_true", help="keep the detector's noise")
    parser.add_argument("--scenes", help="comma-separated scene names (default: all)")
    parser.add_argument("--fresh", action="store_true", help="render commits again")
    args = parser.parse_args()

    if args.base is None and args.head is None:
        args.base, args.head = ("HEAD", "worktree") if dirty() else ("HEAD~1", "HEAD")
    base = Side(args.base or "HEAD", "base")
    head = Side(args.head or "worktree", "head")
    print(f"before: {base.label}\nafter:  {head.label}", flush=True)
    started = time.time()
    for side in (base, head):
        side.prepare()
        side.build()
    noise = args.noise
    if not noise and not (base.has_noiseless and head.has_noiseless):
        print("one revision's CLI has no --noiseless: rendering both with noise", flush=True)
        noise = True

    before_scenes, after_scenes = base.scenes(), head.scenes()
    names = sorted(before_scenes.keys() | after_scenes.keys())
    if args.scenes:
        wanted = set(args.scenes.split(","))
        names = [n for n in names if n in wanted]
    before = base.render_all(names, noise, args.fresh)
    after = head.render_all(names, noise, args.fresh)

    shutil.rmtree(OUT, ignore_errors=True)
    OUT.mkdir(parents=True)
    items = []
    for name in names:
        scene_changed = (
            name in before_scenes
            and name in after_scenes
            and before_scenes[name].read_bytes() != after_scenes[name].read_bytes()
        )
        for tier in TIERS:
            items.append(
                compare(before.get((name, tier)), after.get((name, tier)), scene_changed,
                        name, tier)
            )
    summary = {
        "base": base.label,
        "head": head.label,
        "noise": noise,
        "tiers": [f"{t}/{CAMERA}" for t in TIERS],
        "made": time.strftime("%Y-%m-%d %H:%M"),
        "items": items,
    }
    (OUT / "summary.json").write_text(json.dumps(summary, indent=1))
    counts = {}
    for item in items:
        counts[item["status"]] = counts.get(item["status"], 0) + 1
    print(
        f"{len(items)} images in {time.time() - started:.0f} s: "
        + ", ".join(f"{n} {s}" for s, n in sorted(counts.items()))
    )
    for item in items:
        if item["status"] == "changed":
            print(
                f"  {item['scene']} {item['tier']}: {item['pixels']} pixels "
                f"({100 * item['changed']:.3g} %), "
                f"up to {item['max_dn']} DN"
            )
    print("view with scripts/serve-gallery.py (the Delta tab)")


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