#!/usr/bin/env python3
"""
shortlist.py - Cherry engine one, per-frame judgment (v0.2 prototype)

A dependency-light, fully local culling pipeline. Point it at a folder of photos
and it returns a ranked shortlist plus an HTML contact sheet you can eyeball.

    pip install opencv-python-headless numpy pillow
    python shortlist.py /path/to/photos --mode top --n 50 --genre event --open

What it does, in order:
  1. Ingest        walk folder, read EXIF (time, camera, lens, ISO, focal length)
  2. Technical     sharpness, exposure, clipping, contrast, noise, colorfulness
  3. Subject       face detection -> face size, face sharpness, eyes-open, smile
  4. Composition   rule-of-thirds, horizon tilt, subject-cropped-at-edge penalty
  5. Grouping      near-duplicate / burst clustering (dHash + colour + EXIF time)
  5b. Brackets     exposure ladders kept whole, never collapsed to one frame
  6. Scoring       genre-weighted blend of the above into a single 0-100 score
  7. Selection     'unique' / 'top N' / 'social', diversity-aware (MMR)
  8. Output        CSV + JSON + HTML contact sheet (+ optional file copy)

Deliberately NO deep learning in the core loop: this tier is fast, free, runs
offline, and gets you ~70-80% of the way. Plug a learned scorer in via
--scorer-plugin once you want the last 20% (see PLUGIN CONTRACT at the bottom).
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import shutil
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any

import cv2
import numpy as np
from PIL import Image, ExifTags

# --------------------------------------------------------------------------
# config
# --------------------------------------------------------------------------

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp", ".heic", ".heif"}
RAW_EXTS = {".cr2", ".cr3", ".nef", ".arw", ".raf", ".orf", ".rw2", ".dng"}

WORK_LONG_EDGE = 1024      # analysis resolution — big enough for real sharpness
THUMB_LONG_EDGE = 320      # contact-sheet thumbnails
HASH_SIZE = 8              # dHash grid -> 64-bit hash

EXIF_TAGS = {v: k for k, v in ExifTags.TAGS.items()}


# Genre weight profiles. Each maps a metric name -> weight. Weights are
# normalised at scoring time, so they're readable as relative importance.
# This table IS the product's opinion about "what makes a good photograph"
# for each genre — it's the thing you tune with real photographer feedback.
GENRE_PROFILES: dict[str, dict[str, float]] = {
    "general": {
        "sharpness": 1.00, "subject_sharpness": 0.80, "exposure": 0.90,
        "clipping": 0.60, "contrast": 0.45, "noise": 0.40,
        "colorfulness": 0.30, "composition": 0.55, "face_quality": 0.60,
        "eyes_open": 0.50, "tilt": 0.30, "edge_crop": 0.35,
    },
    # Weddings / events: faces rule. A tack-sharp empty chair loses to a
    # slightly soft frame where everyone's eyes are open.
    "event": {
        "sharpness": 0.55, "subject_sharpness": 1.00, "exposure": 0.85,
        "clipping": 0.55, "contrast": 0.30, "noise": 0.30,
        "colorfulness": 0.20, "composition": 0.40, "face_quality": 1.20,
        "eyes_open": 1.10, "tilt": 0.25, "edge_crop": 0.60,
    },
    "portrait": {
        "sharpness": 0.45, "subject_sharpness": 1.30, "exposure": 0.85,
        "clipping": 0.70, "contrast": 0.30, "noise": 0.45,
        "colorfulness": 0.20, "composition": 0.55, "face_quality": 1.40,
        "eyes_open": 1.30, "tilt": 0.30, "edge_crop": 0.70,
    },
    # Real estate: geometry and dynamic range are everything, faces are noise.
    "realestate": {
        "sharpness": 1.20, "subject_sharpness": 0.30, "exposure": 1.20,
        "clipping": 1.30, "contrast": 0.55, "noise": 0.60,
        "colorfulness": 0.25, "composition": 0.70, "face_quality": 0.00,
        "eyes_open": 0.00, "tilt": 1.30, "edge_crop": 0.10,
    },
    # Food: shallow DOF is intentional, so global sharpness matters less than
    # subject sharpness; colour and clean highlights matter a lot.
    "food": {
        "sharpness": 0.40, "subject_sharpness": 1.20, "exposure": 1.00,
        "clipping": 1.00, "contrast": 0.60, "noise": 0.55,
        "colorfulness": 0.95, "composition": 0.85, "face_quality": 0.00,
        "eyes_open": 0.00, "tilt": 0.60, "edge_crop": 0.15,
    },
    "landscape": {
        "sharpness": 1.40, "subject_sharpness": 0.35, "exposure": 1.00,
        "clipping": 1.10, "contrast": 0.75, "noise": 0.70,
        "colorfulness": 0.85, "composition": 0.90, "face_quality": 0.00,
        "eyes_open": 0.00, "tilt": 1.20, "edge_crop": 0.10,
    },
    # 'social' is a selection mode, but it also shifts the weights: punchy,
    # colourful, face-forward frames that survive a 4-inch screen.
    "social": {
        "sharpness": 0.55, "subject_sharpness": 0.80, "exposure": 0.80,
        "clipping": 0.50, "contrast": 0.90, "noise": 0.25,
        "colorfulness": 1.10, "composition": 0.85, "face_quality": 0.85,
        "eyes_open": 0.80, "tilt": 0.35, "edge_crop": 0.45,
    },
}

METRIC_KEYS = list(GENRE_PROFILES["general"].keys())


# --------------------------------------------------------------------------
# data model
# --------------------------------------------------------------------------

@dataclass
class Face:
    x: int
    y: int
    w: int
    h: int
    confidence: float = 1.0
    sharpness: float = 0.0        # 0-1, relative to frame
    eyes_open: float = 0.5        # 0-1, 0.5 == unknown
    smiling: float = 0.0          # 0-1
    area_frac: float = 0.0        # fraction of frame area


@dataclass
class Shot:
    path: str
    filename: str = ""
    ok: bool = True
    error: str = ""

    width: int = 0
    height: int = 0
    aspect: float = 0.0
    filesize: int = 0

    # EXIF
    timestamp: float = 0.0
    camera: str = ""
    lens: str = ""
    iso: int = 0
    focal_mm: float = 0.0
    aperture: float = 0.0
    shutter: float = 0.0
    exposure_bias: float = 0.0    # EV compensation dialled in on the camera
    ev: float | None = None       # computed exposure value, None if unknowable

    # raw measurements (unnormalised, useful for debugging + retraining)
    raw: dict[str, float] = field(default_factory=dict)
    # normalised 0-1 metrics, higher == better
    m: dict[str, float] = field(default_factory=dict)

    faces: list[Face] = field(default_factory=list)
    faces_checked: bool = False   # False == no detector ran, so "no faces" is
                                  # not evidence of anything

    dhash: int = 0
    hist: list[float] = field(default_factory=list)

    cluster: int = -1
    bracket: int = -1             # -1 == not part of an exposure bracket
    bracket_role: str = ""        # 'under' | 'base' | 'over'
    bracket_size: int = 0
    score: float = 0.0
    selected: bool = False
    reasons: list[str] = field(default_factory=list)
    keep_reasons: list[str] = field(default_factory=list)


# --------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------

def _clamp01(x: float) -> float:
    return 0.0 if x < 0 else (1.0 if x > 1 else float(x))


def _lin(x: float, lo: float, hi: float) -> float:
    """Map x from [lo,hi] onto [0,1], clamped."""
    if hi == lo:
        return 0.0
    return _clamp01((x - lo) / (hi - lo))


def _load_bgr(path: str, long_edge: int = WORK_LONG_EDGE) -> tuple[np.ndarray | None, int, int]:
    """Load and downscale. Returns (bgr, orig_w, orig_h)."""
    ext = Path(path).suffix.lower()

    if ext in RAW_EXTS:
        try:
            import rawpy  # optional
            with rawpy.imread(path) as r:
                rgb = r.postprocess(use_camera_wb=True, half_size=True,
                                    no_auto_bright=False, output_bps=8)
            img = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
            oh, ow = img.shape[:2]
            ow, oh = ow * 2, oh * 2  # half_size
        except Exception:
            return None, 0, 0
    else:
        try:
            if ext in (".heic", ".heif"):
                try:
                    import pillow_heif  # optional
                    pillow_heif.register_heif_opener()
                except Exception:
                    return None, 0, 0
            with Image.open(path) as im:
                ow, oh = im.size
                # libjpeg can decode straight to 1/2, 1/4 or 1/8 scale, far
                # cheaper than decoding 24MP and then resizing. Culling is
                # decode-bound, so this is most of the runtime. draft() picks
                # the largest reduction that still meets the requested size, so
                # asking for exactly `long_edge` never decodes below the
                # resolution the sharpness metrics are calibrated for.
                if max(ow, oh) > long_edge:
                    k = long_edge / max(ow, oh)
                    im.draft("RGB", (max(1, round(ow * k)), max(1, round(oh * k))))
                im = im.convert("RGB")
                img = cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)
        except Exception:
            return None, 0, 0

    h, w = img.shape[:2]
    scale = long_edge / max(h, w)
    if scale < 1.0:
        img = cv2.resize(img, (max(1, int(w * scale)), max(1, int(h * scale))),
                         interpolation=cv2.INTER_AREA)
    return img, ow, oh


def _read_exif(path: str) -> dict[str, Any]:
    out: dict[str, Any] = {}
    try:
        with Image.open(path) as im:
            ex = im.getexif()
            if not ex:
                return out
            # getexif() is IFD0 only. Capture time, aperture, shutter, ISO and
            # focal length live in the Exif sub-IFD behind tag 0x8769, so this
            # has to be fetched separately or every one of them reads empty.
            try:
                exif_ifd = ex.get_ifd(0x8769)
            except Exception:
                exif_ifd = {}

            def g(name):
                tid = EXIF_TAGS.get(name)
                if tid is None:
                    return None
                v = exif_ifd.get(tid)
                return ex.get(tid) if v is None else v

            dt = g("DateTimeOriginal") or g("DateTime")
            if dt:
                try:
                    out["timestamp"] = time.mktime(time.strptime(str(dt), "%Y:%m:%d %H:%M:%S"))
                except Exception:
                    pass
            make, model = g("Make"), g("Model")
            if make or model:
                out["camera"] = f"{make or ''} {model or ''}".strip()
            if g("LensModel"):
                out["lens"] = str(g("LensModel"))
            for src, dst, cast in (("ISOSpeedRatings", "iso", int),
                                   ("FocalLength", "focal_mm", float),
                                   ("FNumber", "aperture", float),
                                   ("ExposureTime", "shutter", float),
                                   ("ExposureBiasValue", "exposure_bias", float)):
                v = g(src)
                if v is not None:
                    try:
                        if isinstance(v, tuple):
                            v = (v[0] / v[1]) if (len(v) == 2 and v[1]) else v[0]
                        out[dst] = cast(v)
                    except Exception:
                        pass
    except Exception:
        pass
    return out


# --------------------------------------------------------------------------
# stage 2 — technical quality
# --------------------------------------------------------------------------

def sharpness_map(gray: np.ndarray, grid: int = 12) -> tuple[float, np.ndarray, float]:
    """Noise-invariant sharpness. Returns (sharpness_signal, tiles, noise_floor).

    Plain variance-of-Laplacian is the textbook blur metric and it is badly
    broken for culling: grain is high-frequency energy, so a noisy ISO-12800
    frame measures as *sharper* than a clean one and high-ISO junk floats to the
    top of the shortlist. Subtracting an analytic white-noise term doesn't fix
    it either, because after demosaicing and JPEG the noise is spatially
    correlated and the analytic estimate is off by an order of magnitude.

    What actually works: sensor noise raises the Laplacian energy of *every*
    tile by roughly the same amount, while real subject detail only raises the
    tiles containing the subject. So take the spread instead of the level —

        sharpness = p90(tiles) - p20(tiles)

    p20 is a flat region (sky, wall, bokeh) and measures the noise floor; p90 is
    the sharpest part of the subject. Their difference is the detail that noise
    can't fake. Measured on the test fixture, adding sigma=26 noise moves this
    number by ~1%, versus ~13x for raw Laplacian variance.

    The tile map also lets us tell 'blurry photo' apart from 'sharp subject,
    intentionally blurred background' — the latter is a *good* photo, and any
    metric that averages over the whole frame throws it away.
    """
    lap = cv2.Laplacian(gray, cv2.CV_32F, ksize=3)
    h, w = gray.shape
    tiles = np.zeros((grid, grid), np.float32)
    for i in range(grid):
        for j in range(grid):
            sub = lap[i * h // grid:(i + 1) * h // grid, j * w // grid:(j + 1) * w // grid]
            tiles[i, j] = float(sub.var())

    floor = float(np.percentile(tiles, 20))
    peak = float(np.percentile(tiles, 90))
    signal = max(0.0, peak - floor)
    return signal, np.maximum(tiles - floor, 0.0), floor


def exposure_stats(gray: np.ndarray) -> dict[str, float]:
    hist = cv2.calcHist([gray], [0], None, [256], [0, 256]).ravel()
    total = hist.sum() or 1.0
    p = hist / total
    mean = float((np.arange(256) * p).sum())
    clipped_hi = float(p[250:].sum())
    clipped_lo = float(p[:6].sum())
    # entropy of the tone distribution — flat, well-spread histograms score high
    nz = p[p > 0]
    entropy = float(-(nz * np.log2(nz)).sum()) / 8.0
    return {
        "mean_luma": mean,
        "clipped_hi": clipped_hi,
        "clipped_lo": clipped_lo,
        "tone_entropy": entropy,
        "contrast_std": float(gray.std()),
    }


def noise_estimate(gray: np.ndarray, grid: int = 16) -> float:
    """Noise sigma measured in the flattest regions of the frame.

    Immerkaer's classic whole-image estimator is content-biased — fur, foliage
    and fabric read as noise. Instead, high-pass the frame, tile it, and take a
    low percentile of the per-tile standard deviations: the quietest tiles are
    sky/wall/bokeh, where whatever high-frequency energy remains really is
    sensor noise. On the test fixture this separates clean from noisy frames by
    ~17x, versus ~5x for Immerkaer.
    """
    g = gray.astype(np.float32)
    hp = g - cv2.GaussianBlur(g, (0, 0), 1.2)
    h, w = gray.shape
    stds = [float(hp[i * h // grid:(i + 1) * h // grid,
                     j * w // grid:(j + 1) * w // grid].std())
            for i in range(grid) for j in range(grid)]
    return float(np.percentile(stds, 15))


def colorfulness(bgr: np.ndarray) -> float:
    """Hasler & Süsstrunk (2003) colourfulness metric."""
    b, g, r = cv2.split(bgr.astype(np.float32))
    rg = np.abs(r - g)
    yb = np.abs(0.5 * (r + g) - b)
    return float(math.sqrt(rg.std() ** 2 + yb.std() ** 2) + 0.3 * math.sqrt(rg.mean() ** 2 + yb.mean() ** 2))


def horizon_tilt(gray: np.ndarray) -> tuple[float, bool]:
    """Estimate frame tilt in degrees from dominant near-axis lines.

    Returns (degrees, confident). Most photographs have no true horizon, and a
    tilt estimate built from three stray edges in a coffee cup is worse than no
    estimate at all — so we return confident=False and let the caller drop the
    metric entirely rather than inventing a penalty.
    """
    # Tilt is a coarse, global property — half resolution costs nothing in
    # accuracy and HoughLinesP is one of the two hot spots in the pipeline.
    small = cv2.resize(gray, (gray.shape[1] // 2, gray.shape[0] // 2),
                       interpolation=cv2.INTER_AREA) if min(gray.shape) > 400 else gray
    edges = cv2.Canny(small, 60, 180, apertureSize=3)
    minlen = int(0.30 * max(small.shape))
    lines = cv2.HoughLinesP(edges, 1, np.pi / 720, threshold=60,
                            minLineLength=minlen, maxLineGap=6)
    if lines is None or len(lines) < 3:
        return 0.0, False

    # OpenCV 4 hands back (N, 1, 4); OpenCV 5 hands back (N, 4). Reshape rather
    # than index, so this reads the same on either.
    segments = np.asarray(lines).reshape(-1, 4)

    devs: list[float] = []
    lens: list[float] = []
    for x1, y1, x2, y2 in segments:
        ang = math.degrees(math.atan2(y2 - y1, x2 - x1)) % 180.0
        d_h = min(ang, 180.0 - ang)               # distance to horizontal
        d_v = abs(ang - 90.0)                     # distance to vertical
        d = min(d_h, d_v)
        if d < 9.0:                               # near-axis lines only
            devs.append(d)
            lens.append(math.hypot(x2 - x1, y2 - y1))

    # A curved coffee-cup rim also produces "lines" a few degrees off vertical.
    # Demand a lot of evidence: many near-axis lines, forming most of the long
    # lines in the frame, all agreeing closely. Anything less is organic
    # content with no horizon to be level with, and we say so.
    if len(devs) < 6 or len(devs) < 0.5 * len(lines):
        return 0.0, False

    devs_a, lens_a = np.array(devs), np.array(lens)
    order = np.argsort(-lens_a)[:16]              # the longest lines carry the signal
    devs_a, lens_a = devs_a[order], lens_a[order]
    if float(devs_a.std()) > 2.0:                 # lines disagree -> texture, not structure
        return 0.0, False
    est = float((devs_a * lens_a).sum() / (lens_a.sum() + 1e-9))
    return est, True


# --------------------------------------------------------------------------
# stage 3 — faces
# --------------------------------------------------------------------------

_FACE_CASCADE = None
_EYE_CASCADE = None
_SMILE_CASCADE = None
_YUNET = None


_CASCADES_TRIED = False


def _get_detectors():
    """Haar cascades, or (None, None, None) where they are unavailable.

    OpenCV 5.0 dropped CascadeClassifier from the main module and stopped
    shipping the cascade XML files, so the 02 August build crashed outright on
    a current install. Returning None here is not a silent downgrade: with no
    detector the face metrics never enter s.m at all, score_shot renormalises
    over what it could actually measure, and the report shows n/a rather than a
    zero that would read as 'no good faces in this frame'.
    """
    global _FACE_CASCADE, _EYE_CASCADE, _SMILE_CASCADE, _CASCADES_TRIED
    if _CASCADES_TRIED:
        return _FACE_CASCADE, _EYE_CASCADE, _SMILE_CASCADE
    _CASCADES_TRIED = True
    try:
        d = cv2.data.haarcascades
        face = cv2.CascadeClassifier(d + "haarcascade_frontalface_default.xml")
        if face.empty():
            raise RuntimeError("cascade files not installed")
        _FACE_CASCADE = face
        _EYE_CASCADE = cv2.CascadeClassifier(d + "haarcascade_eye_tree_eyeglasses.xml")
        _SMILE_CASCADE = cv2.CascadeClassifier(d + "haarcascade_smile.xml")
    except Exception:
        _FACE_CASCADE = _EYE_CASCADE = _SMILE_CASCADE = None
    return _FACE_CASCADE, _EYE_CASCADE, _SMILE_CASCADE


def _get_yunet(size: tuple[int, int]):
    """YuNet is far better than Haar. Used if the model file is present.

    Download once:
      curl -LO https://raw.githubusercontent.com/opencv/opencv_zoo/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx
    and place it next to this script (or set YUNET_MODEL).
    """
    global _YUNET
    model = os.environ.get("YUNET_MODEL") or str(Path(__file__).parent / "face_detection_yunet_2023mar.onnx")
    if not Path(model).exists():
        return None
    try:
        if _YUNET is None:
            _YUNET = cv2.FaceDetectorYN.create(model, "", size, 0.7, 0.3, 5000)
        _YUNET.setInputSize(size)
        return _YUNET
    except Exception:
        return None


def face_detection_available() -> bool:
    """True when some face detector can actually run on this machine."""
    model = os.environ.get("YUNET_MODEL") or str(Path(__file__).parent / "face_detection_yunet_2023mar.onnx")
    if Path(model).exists():
        return True
    return _get_detectors()[0] is not None


def detect_faces(bgr: np.ndarray, gray: np.ndarray, tiles: np.ndarray,
                 noise_floor: float = 0.0) -> list[Face]:
    h, w = gray.shape
    boxes: list[tuple[int, int, int, int, float]] = []

    yn = _get_yunet((w, h))
    if yn is not None:
        try:
            _, dets = yn.detect(bgr)
            if dets is not None:
                for d in dets:
                    x, y, fw, fh = [int(v) for v in d[:4]]
                    boxes.append((x, y, fw, fh, float(d[-1])))
        except Exception:
            pass

    face_c, eye_c, smile_c = _get_detectors()

    if not boxes and face_c is not None:
        # Detect on a half-size copy (the cascade is the pipeline's other hot
        # spot) and scale the boxes back up - face sharpness is still measured
        # on the full working-resolution pixels.
        ds = 2 if min(h, w) > 400 else 1
        eq = cv2.equalizeHist(gray if ds == 1 else
                              cv2.resize(gray, (w // ds, h // ds), interpolation=cv2.INTER_AREA))
        found = face_c.detectMultiScale(eq, scaleFactor=1.12, minNeighbors=6,
                                        minSize=(max(20, w // (40 * ds)),) * 2)
        boxes = [(int(x * ds), int(y * ds), int(fw * ds), int(fh * ds), 0.8)
                 for x, y, fw, fh in found]

    frame_area = float(w * h)
    faces: list[Face] = []

    for x, y, fw, fh, conf in boxes:
        x, y = max(0, x), max(0, y)
        fw, fh = min(fw, w - x), min(fh, h - y)
        if fw < 12 or fh < 12:
            continue
        roi = gray[y:y + fh, x:x + fw]
        if roi.size == 0:
            continue

        # Measure face sharpness *within the face*, never against the frame.
        # Skin is intrinsically smoother than fabric or foliage, so a portrait
        # in front of a detailed background scores near zero on any
        # frame-relative measure — every sharp portrait gets flagged soft. The
        # same p90-p20 spread trick works here: eyelashes and lip edges are the
        # p90, cheeks are the p20, and their gap is what focus actually buys.
        if min(roi.shape) >= 24:
            f_sig, _, _ = sharpness_map(roi, grid=6)
        else:
            f_sig = max(0.0, float(cv2.Laplacian(roi, cv2.CV_32F).var()) - noise_floor)
        f_sharp = _lin(math.log1p(f_sig), math.log1p(20), math.log1p(1500))
        # eyes: upper 60% of the face box
        eye_roi = roi[: int(fh * 0.6), :]
        if eye_c is not None:
            eyes = eye_c.detectMultiScale(cv2.equalizeHist(eye_roi), 1.1, 5,
                                          minSize=(max(6, fw // 12), max(6, fh // 16)))
            eyes_open = 0.5 if len(eyes) == 0 else min(1.0, 0.55 + 0.25 * len(eyes))
        else:
            eyes_open = 0.5      # 0.5 is this file's value for "not knowable"
        # smile: lower 45%
        sm_roi = roi[int(fh * 0.55):, :]
        if smile_c is not None and sm_roi.size:
            smiles = smile_c.detectMultiScale(cv2.equalizeHist(sm_roi), 1.7, 20)
            smiling = 1.0 if len(smiles) else 0.0
        else:
            smiling = 0.0

        faces.append(Face(
            x=x, y=y, w=fw, h=fh, confidence=conf,
            sharpness=_clamp01(f_sharp),
            eyes_open=eyes_open, smiling=smiling,
            area_frac=(fw * fh) / frame_area,
        ))

    faces.sort(key=lambda f: -f.area_frac)
    return faces[:20]


# --------------------------------------------------------------------------
# stage 4 — composition
# --------------------------------------------------------------------------

def subject_centroid(tiles: np.ndarray, faces: list[Face], shape: tuple[int, int]) -> tuple[float, float]:
    """Where the 'subject' is, in normalised 0-1 coords."""
    h, w = shape
    if faces:
        big = faces[0]
        return ((big.x + big.w / 2) / w, (big.y + big.h / 2) / h)
    # fall back to the detail centre of mass
    t = tiles / (tiles.sum() + 1e-9)
    ty, tx = t.shape
    ys, xs = np.mgrid[0:ty, 0:tx]
    cy = float((t * (ys + 0.5) / ty).sum())
    cx = float((t * (xs + 0.5) / tx).sum())
    return cx, cy


def thirds_score(cx: float, cy: float) -> float:
    """1.0 at a rule-of-thirds intersection or dead centre, falling off with distance.

    Dead centre is included on purpose — centred symmetry is a legitimate,
    common composition, and penalising it produces obviously wrong culls.
    """
    anchors = [(1 / 3, 1 / 3), (2 / 3, 1 / 3), (1 / 3, 2 / 3), (2 / 3, 2 / 3), (0.5, 0.5)]
    d = min(math.hypot(cx - ax, cy - ay) for ax, ay in anchors)
    return _clamp01(1.0 - d / 0.30)


def edge_crop_penalty(faces: list[Face], shape: tuple[int, int]) -> float:
    """0 == no faces clipped by the frame edge, 1 == badly clipped."""
    if not faces:
        return 0.0
    h, w = shape
    worst = 0.0
    for f in faces[:5]:
        if f.area_frac < 0.005:
            continue
        margin = min(f.x, f.y, w - (f.x + f.w), h - (f.y + f.h))
        if margin < 0:
            worst = 1.0
        elif margin < 0.02 * min(w, h):
            worst = max(worst, 0.6)
    return worst


# --------------------------------------------------------------------------
# stage 5 — near-duplicate / burst grouping
# --------------------------------------------------------------------------

def dhash(gray: np.ndarray, size: int = HASH_SIZE) -> int:
    small = cv2.resize(gray, (size + 1, size), interpolation=cv2.INTER_AREA)
    diff = small[:, 1:] > small[:, :-1]
    v = 0
    for bit in diff.flatten():
        v = (v << 1) | int(bit)
    return v


def color_hist(bgr: np.ndarray) -> list[float]:
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    h = cv2.calcHist([hsv], [0, 1], None, [12, 6], [0, 180, 0, 256])
    h = cv2.normalize(h, h).flatten()
    return [float(x) for x in h]


def hamming(a: int, b: int) -> int:
    return bin(a ^ b).count("1")


def cluster_shots(shots: list[Shot], hash_thresh: int = 12,
                  hist_thresh: float = 0.55, time_gap: float = 8.0) -> int:
    """Union-find over a sliding window. Two frames join a cluster if they're
    perceptually close AND (shot close in time OR very perceptually close).

    Sliding window keeps this O(n·k) instead of O(n²) — a 5000-image wedding
    doesn't need all-pairs comparison because bursts are contiguous in time.
    """
    order = sorted(range(len(shots)), key=lambda i: (shots[i].timestamp or 0, shots[i].filename))
    parent = list(range(len(shots)))

    def find(a):
        while parent[a] != a:
            parent[a] = parent[parent[a]]
            a = parent[a]
        return a

    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra

    WINDOW = 25
    for oi in range(len(order)):
        i = order[oi]
        if not shots[i].ok:
            continue
        hi = np.array(shots[i].hist, np.float32)
        for oj in range(oi + 1, min(oi + 1 + WINDOW, len(order))):
            j = order[oj]
            if not shots[j].ok:
                continue
            hd = hamming(shots[i].dhash, shots[j].dhash)
            if hd > hash_thresh + 8:
                continue
            hj = np.array(shots[j].hist, np.float32)
            corr = float(cv2.compareHist(hi, hj, cv2.HISTCMP_CORREL))
            dt = abs((shots[i].timestamp or 0) - (shots[j].timestamp or 0))
            close_in_time = (shots[i].timestamp and shots[j].timestamp and dt <= time_gap)
            if hd <= hash_thresh and corr >= hist_thresh and (close_in_time or hd <= 6):
                union(i, j)

    roots: dict[int, int] = {}
    for i in range(len(shots)):
        r = find(i)
        if r not in roots:
            roots[r] = len(roots)
        shots[i].cluster = roots[r]
    return len(roots)


# --------------------------------------------------------------------------
# stage 5b - exposure brackets
# --------------------------------------------------------------------------

def _ev(s: Shot) -> float | None:
    """Exposure value from the camera's own numbers, or None if unknowable.

    Higher EV means less light reached the sensor, so a high-EV frame is the
    dark end of a ladder. Nothing here is inferred from pixels: an accidental
    two-stop overexposure and a deliberate +2 EV bracket frame look identical
    in the image, and are only told apart by the metadata.
    """
    if not s.aperture or not s.shutter or s.shutter <= 0:
        return None
    return math.log2((s.aperture * s.aperture) / s.shutter)


def detect_brackets(shots: list[Shot], time_gap: float = 3.0,
                    hash_thresh: int = 12, min_frames: int = 3,
                    min_span_stops: float = 1.0) -> int:
    """Find exposure ladders and mark them as one logical image.

    Real estate, interiors and architecture are shot as brackets: the camera
    stays on a tripod and fires three to seven frames of the same framing at
    different exposures, to be merged later. Every culling tool in this market
    reads that as a burst and returns one frame, which destroys the shot. This
    file did the same until this function existed.

    A run qualifies only when the camera's own metadata proves it:
      - at least `min_frames` frames, consecutive in time
      - each gap under `time_gap` seconds
      - framing effectively identical (perceptual hash within `hash_thresh`;
        two stops of exposure move a real frame's hash by about 8 bits, so this
        is looser than it looks and the metadata below does the real work)
      - one focal length and one ISO across the run, so it is not a reframe
      - three or more distinct exposures spanning at least `min_span_stops`

    Returns the number of brackets found. Frames in a bracket keep their normal
    scores for sharpness and composition; what changes is that exposure and
    clipping stop counting against them, and selection treats the ladder as a
    single image (see score_shot and select).
    """
    for s in shots:
        s.bracket, s.bracket_role, s.bracket_size = -1, "", 0
        s.ev = _ev(s)

    usable = [s for s in shots if s.ok and s.timestamp and s.ev is not None]
    usable.sort(key=lambda s: (s.timestamp, s.filename))

    n_found = 0
    i = 0
    while i < len(usable):
        run = [usable[i]]
        j = i + 1
        while j < len(usable):
            prev, cand = run[-1], usable[j]
            same_rig = (
                abs(cand.timestamp - prev.timestamp) <= time_gap
                and hamming(run[0].dhash, cand.dhash) <= hash_thresh
                and (not run[0].focal_mm or not cand.focal_mm
                     or abs(run[0].focal_mm - cand.focal_mm) <= 0.02 * run[0].focal_mm)
                and run[0].iso == cand.iso
            )
            if not same_rig:
                break
            run.append(cand)
            j += 1

        if len(run) >= min_frames:
            # Brightness order. Exposure compensation is the clearest evidence
            # when the camera recorded it; computed EV is the fallback.
            biases = set(round(s.exposure_bias, 2) for s in run)
            if len(biases) >= min_frames:
                key = lambda s: s.exposure_bias
            else:
                key = lambda s: -(s.ev or 0.0)
            levels = sorted(set(round(key(s), 2) for s in run))
            if len(levels) >= min_frames and (levels[-1] - levels[0]) >= min_span_stops:
                ladder = sorted(run, key=key)
                mid = ladder[len(ladder) // 2]
                for s in ladder:
                    s.bracket = n_found
                    s.bracket_size = len(ladder)
                    s.bracket_role = "base" if s is mid else (
                        "under" if key(s) < key(mid) else "over")
                n_found += 1
                i = j
                continue
        i += 1
    return n_found


def bracket_groups(shots: list[Shot]) -> dict[int, list[Shot]]:
    """Bracket id -> its frames, best-scoring first."""
    out: dict[int, list[Shot]] = {}
    for s in shots:
        if s.bracket >= 0:
            out.setdefault(s.bracket, []).append(s)
    for v in out.values():
        v.sort(key=lambda s: -s.score)
    return out


# --------------------------------------------------------------------------
# per-image worker
# --------------------------------------------------------------------------

def analyse(path: str) -> Shot:
    s = Shot(path=path, filename=Path(path).name)
    try:
        s.filesize = os.path.getsize(path)
    except OSError:
        pass

    bgr, ow, oh = _load_bgr(path)
    if bgr is None:
        s.ok, s.error = False, "could not decode"
        return s

    s.width, s.height = ow, oh
    s.aspect = (ow / oh) if oh else 0.0
    for k, v in _read_exif(path).items():
        setattr(s, k, v)

    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    h, w = gray.shape

    noise = noise_estimate(gray)
    g_sharp, tiles, noise_floor = sharpness_map(gray)
    ex = exposure_stats(gray)
    cf = colorfulness(bgr)
    tilt, tilt_ok = horizon_tilt(gray)
    s.faces_checked = face_detection_available()
    faces = detect_faces(bgr, gray, tiles, noise_floor)
    cx, cy = subject_centroid(tiles, faces, (h, w))

    # Best tile = the sharpest region of the frame. If there's a face, we care
    # about the sharpness *there* — a tack-sharp background behind a soft face
    # is a reject, and the whole-frame number can't tell you that.
    tile_best = float(tiles.max())
    tile_med = float(np.median(tiles))

    # Face sharpness is already normalised 0-1 on its own scale; the frame-level
    # signal is a raw variance. Keep them apart and only convert at the end.
    face_sharp = max((f.sharpness for f in faces[:3]), default=None)

    s.raw = {
        "sharp_global": g_sharp, "tile_best": tile_best, "tile_median": tile_med,
        "face_sharp": -1.0 if face_sharp is None else face_sharp,
        "noise_sigma": noise, "colorfulness": cf,
        "tilt_deg": tilt, "tilt_confident": float(tilt_ok), "n_faces": float(len(faces)),
        "largest_face_frac": faces[0].area_frac if faces else 0.0,
        "cx": cx, "cy": cy, **ex,
    }

    # ---- normalise to 0-1, higher is better -----------------------------
    # Thresholds below are calibrated for WORK_LONG_EDGE=1024, 8-bit sRGB.
    # These are the numbers you'd replace with a learned model.
    # A value of None means "not measurable on this frame" — the metric is
    # dropped from the weighted average instead of guessed at.
    m: dict[str, float | None] = {}
    m["sharpness"] = _lin(math.log1p(g_sharp), math.log1p(30), math.log1p(2000))
    # If there's a face, that IS the subject and its focus is what matters.
    # Otherwise fall back to the sharpest region of the frame.
    m["subject_sharpness"] = (face_sharp if face_sharp is not None
                              else _lin(math.log1p(tile_best), math.log1p(30), math.log1p(2200)))

    # Exposure as a plateau, not a target. Plenty of good photographs are low-key
    # or high-key on purpose; only punish luma once it leaves the range where
    # detail actually survives. Highlight clipping is weighted ~2x shadow
    # clipping because blown highlights are unrecoverable and crushed blacks
    # usually aren't.
    luma = ex["mean_luma"]
    off = 0.0 if 72.0 <= luma <= 168.0 else (72.0 - luma if luma < 72.0 else luma - 168.0)
    m["exposure"] = _clamp01(1.0 - off / 55.0) * (0.55 + 0.45 * ex["tone_entropy"])
    m["clipping"] = _clamp01(1.0 - (ex["clipped_hi"] * 9.0 + ex["clipped_lo"] * 4.0))
    m["contrast"] = _lin(ex["contrast_std"], 18, 72)
    m["noise"] = 1.0 - _lin(noise, 1.5, 7.0)
    m["colorfulness"] = _lin(cf, 8, 62)
    m["composition"] = thirds_score(cx, cy)
    m["tilt"] = (1.0 - _lin(tilt, 0.7, 5.0)) if tilt_ok else None
    m["edge_crop"] = (1.0 - edge_crop_penalty(faces, (h, w))) if faces else None

    if faces:
        top = faces[:4]
        size_term = _lin(math.sqrt(faces[0].area_frac), 0.05, 0.40)
        sharp_term = float(np.mean([f.sharpness for f in top]))
        smile_term = float(np.mean([f.smiling for f in top]))
        m["face_quality"] = _clamp01(0.45 * size_term + 0.40 * sharp_term + 0.15 * smile_term)
        m["eyes_open"] = float(np.mean([f.eyes_open for f in top]))
    else:
        # No face is not a fault — a landscape shouldn't be punished for having
        # no people in it. Drop the face metrics entirely for this frame.
        m["face_quality"] = None
        m["eyes_open"] = None

    s.m = {k: (None if v is None else round(_clamp01(v), 4)) for k, v in m.items()}
    s.faces = faces
    s.dhash = dhash(gray)
    s.hist = color_hist(bgr)
    return s


# --------------------------------------------------------------------------
# scoring + selection
# --------------------------------------------------------------------------

def score_shot(s: Shot, genre: str, plugin_scores: dict[str, float] | None = None) -> float:
    prof = GENRE_PROFILES.get(genre, GENRE_PROFILES["general"])

    # Only metrics that were actually measurable on this frame contribute, and
    # the weights renormalise over that subset. This is why a landscape with no
    # faces isn't dragged to the middle of the pack by two missing metrics.
    # A frame inside an exposure bracket is *meant* to be dark or blown: that
    # is the job of the -2 and +2 frames. Judging them on exposure or clipping
    # would throw away the two frames the merge needs most, so those metrics
    # drop out of the blend and out of the gates below.
    muted = {"exposure", "clipping"} if s.bracket >= 0 else set()

    num = den = 0.0
    for k, wgt in prof.items():
        v = s.m.get(k)
        if v is None or wgt == 0.0 or k in muted:
            continue
        num += wgt * v
        den += wgt
    base = (num / den) if den else 0.0

    # Hard gates. Some faults are disqualifying no matter how good the rest of
    # the frame is — a photographer will never deliver an out-of-focus frame
    # because it was nicely composed.
    penalties: list[str] = []

    def gate(key: str, thresh: float, mult: float, label: str) -> None:
        nonlocal base
        if key in muted:
            return
        v = s.m.get(key)
        if v is not None and v < thresh:
            base *= mult
            penalties.append(label)

    # Focus gate looks at the best evidence of focus anywhere in the frame:
    # subject sharpness alone would clear a defocused portrait whose face still
    # scores mid-range on its own local scale, and global sharpness alone would
    # condemn a legitimate shallow-depth-of-field shot.
    # subject_sharpness is measured on the face when a detector found one, and
    # estimated from the sharpest tiles when it did not. The estimate is a fair
    # positive signal and a poor veto: on a defocused portrait it reads around
    # 0.4 and would lift the frame straight over this gate. So it only counts
    # here when a detector actually ran and found a subject to measure.
    subject = s.m.get("subject_sharpness") if (s.faces_checked and s.faces) else None
    focus = max(s.m.get("sharpness") or 0.0, subject or 0.0)
    if focus < 0.22:
        base *= 0.35
        penalties.append("very soft")

    gate("clipping", 0.25, 0.70, "clipped")
    gate("clipping", 0.08, 0.70, "badly blown")     # compounds with the above
    gate("exposure", 0.18, 0.65, "badly exposed")
    gate("eyes_open", 0.45, 0.75, "eyes likely closed")
    gate("noise", 0.15, 0.80, "very noisy")

    if plugin_scores and s.path in plugin_scores:
        base = 0.6 * base + 0.4 * _clamp01(plugin_scores[s.path])
        penalties.append("plugin-scored")

    if s.bracket >= 0:
        penalties.append("bracket %s of %d" % (s.bracket_role, s.bracket_size))

    s.reasons = penalties
    return round(100.0 * _clamp01(base), 2)


# Human labels for the metrics a keep-reason can cite. Higher is always
# better in s.m, so "clean" means the noise metric is high, not low.
_KEEP_LABELS = {
    "sharpness": "sharp",
    "subject_sharpness": "subject in focus",
    "exposure": "well exposed",
    "clipping": "highlights held",
    "contrast": "good contrast",
    "noise": "clean",
    "colorfulness": "strong colour",
    "composition": "well composed",
    "tilt": "level",
    "edge_crop": "nothing cropped at the edge",
    "face_quality": "face sharp",
    "eyes_open": "eyes open",
}


def explain_keeps(shots: list[Shot], genre: str) -> None:
    """Write a keep_reasons list onto every selected shot.

    A cut already explains itself through the penalty gates. A keep never did:
    the report showed PICK and a bare score, which is exactly the "score with
    no reason" habit the product refuses. The reason is built from the same
    facts the selection used - what the frame beat inside its own moment, and
    which profile-weighted metrics carried it - so it cannot drift from the
    decision it explains. When nothing about the frame is strong, the honest
    reason is that it was the best of a weak moment, and it says so.
    """
    prof = GENRE_PROFILES.get(genre, GENRE_PROFILES["general"])
    by_cluster: dict[int, list[Shot]] = {}
    for s in shots:
        if s.ok:
            by_cluster.setdefault(s.cluster, []).append(s)

    for s in shots:
        if not s.selected:
            continue
        why: list[str] = []

        mates = by_cluster.get(s.cluster, [s])
        rivals = [m for m in mates if m is not s]
        if s.bracket >= 0:
            why.append("bracket ladder of %d, competes as one photo" % s.bracket_size)
            # the other rungs of its own ladder are companions, not rivals
            rivals = [m for m in rivals if m.bracket != s.bracket]
        if rivals:
            best_rival = max(rivals, key=lambda m: m.score)
            if s.score >= best_rival.score:
                why.append("best of %d in its moment, next scored %.0f"
                           % (len(mates), best_rival.score))
            elif best_rival.selected:
                why.append("second keeper from a strong moment, %.1f behind the first"
                           % (best_rival.score - s.score))
            else:
                why.append("kept for spread: its moment was otherwise unrepresented")

        # The two strongest profile-weighted metrics, only if genuinely strong.
        muted = {"exposure", "clipping"} if s.bracket >= 0 else set()
        strong = sorted(
            ((prof.get(k, 0.0) * v, k) for k, v in s.m.items()
             if v is not None and v >= 0.72 and prof.get(k, 0.0) > 0 and k not in muted),
            reverse=True)
        for _, k in strong[:2]:
            lbl = _KEEP_LABELS.get(k)
            if lbl:
                why.append(lbl)

        if len(why) <= 1 and not strong:
            why.append("nothing about it is strong; least bad in a weak moment")

        s.keep_reasons = why


def _with_bracket_mates(picked: list[Shot],
                        companions: dict[int, list[Shot]]) -> list[Shot]:
    """Put the rest of each selected bracket back, immediately after its base."""
    if not companions:
        return picked
    out: list[Shot] = []
    for s in picked:
        out.append(s)
        if s.bracket >= 0:
            out.extend(companions.get(s.bracket, []))
    return out


def select(shots: list[Shot], mode: str, n: int, per_cluster: int = 1,
           floor: float = 40.0, keep_brackets: bool = True) -> list[Shot]:
    """Diversity-aware selection.

    'unique'  -> best `per_cluster` from each near-duplicate cluster
    'top'     -> one pass of cluster winners, then fill by score, capped so no
                 single cluster can eat the shortlist
    'social'  -> like 'top' but on the social-weighted score, tighter cap

    `floor` stops "one best frame per moment" from smuggling garbage in: if a
    whole burst is out of focus, there is no keeper in it, and the honest answer
    is to return fewer than N rather than pad the shortlist. The floor is only
    relaxed if it would leave the shortlist empty.
    """
    live = [s for s in shots if s.ok]
    if not live:
        return []

    # An exposure bracket is one photograph delivered as several files. It
    # competes for one slot, and if it wins the whole ladder goes with it -
    # handing back a single frame of a bracket is handing back nothing.
    companions: dict[int, list[Shot]] = {}
    if keep_brackets:
        groups = bracket_groups(live)
        for bid, members in groups.items():
            rep = next((m for m in members if m.bracket_role == "base"), members[0])
            companions[bid] = [m for m in members if m is not rep]
        hidden = set(id(m) for ms in companions.values() for m in ms)
        if hidden:
            live = [s for s in live if id(s) not in hidden]

    by_cluster: dict[int, list[Shot]] = {}
    for s in live:
        by_cluster.setdefault(s.cluster, []).append(s)
    for v in by_cluster.values():
        v.sort(key=lambda s: -s.score)

    # relax the floor rather than return nothing at all
    if not any(s.score >= floor for s in live):
        floor = max(s.score for s in live) * 0.75

    if mode == "unique":
        out = [s for v in by_cluster.values() for s in v[:per_cluster] if s.score >= floor]
        out.sort(key=lambda s: -s.score)
        return _with_bracket_mates(out, companions)

    cap = 1 if mode == "social" else max(1, per_cluster)
    picked: list[Shot] = []
    picked_ids: set[int] = set()
    taken: dict[int, int] = {}

    # round 1: the best frame from every distinct moment, if it clears the bar
    for v in sorted(by_cluster.values(), key=lambda v: -v[0].score):
        if len(picked) >= n:
            break
        if v[0].score < floor:
            continue
        picked.append(v[0])
        picked_ids.add(id(v[0]))
        taken[v[0].cluster] = 1

    # Round 2+: fill remaining slots by score. The per-cluster cap starts tight
    # so a single 40-frame burst can't eat the shortlist, then relaxes one step
    # at a time if we still can't fill N. The floor is never relaxed — running
    # out of *good* frames means returning fewer, not returning bad ones.
    rest = sorted((s for s in live if id(s) not in picked_ids and s.score >= floor),
                  key=lambda s: -s.score)
    limit = cap + 1
    while len(picked) < n and any(id(s) not in picked_ids for s in rest):
        for s in rest:
            if len(picked) >= n:
                break
            if id(s) in picked_ids or taken.get(s.cluster, 0) >= limit:
                continue
            picked.append(s)
            picked_ids.add(id(s))
            taken[s.cluster] = taken.get(s.cluster, 0) + 1
        limit += 1

    picked.sort(key=lambda s: -s.score)
    return _with_bracket_mates(picked[:n], companions)


# --------------------------------------------------------------------------
# output
# --------------------------------------------------------------------------

def write_thumb(path: str, dest: Path) -> bool:
    img, _, _ = _load_bgr(path, THUMB_LONG_EDGE)
    if img is None:
        return False
    return bool(cv2.imwrite(str(dest), img, [cv2.IMWRITE_JPEG_QUALITY, 78]))


def write_html(shots: list[Shot], picked: list[Shot], outdir: Path,
               thumbs: dict[str, str], args) -> Path:
    pick_set = {s.path for s in picked}
    ordered = sorted((s for s in shots if s.ok), key=lambda s: -s.score)
    n_ok = len(ordered)
    n_clusters = len({s.cluster for s in ordered})

    def card(s: Shot) -> str:
        t = thumbs.get(s.path, "")
        sel = "sel" if s.path in pick_set else ""
        def bar(k: str) -> str:
            v = s.m.get(k)
            if v is None:
                return (f'<div class="bar na"><span>{k.replace("_", " ")}</span>'
                        f'<i style="width:0%"></i><b>n/a</b></div>')
            return (f'<div class="bar"><span>{k.replace("_", " ")}</span>'
                    f'<i style="width:{v * 100:.0f}%"></i><b>{v:.2f}</b></div>')
        bars = "".join(bar(k) for k in METRIC_KEYS)
        flags = list(s.reasons)
        if s.path in pick_set and s.keep_reasons:
            flags = [f"kept: {r}" for r in s.keep_reasons] + flags
        why = ("<div class='flags'>" + " ".join(f"<em>{r}</em>" for r in flags) + "</div>") if flags else ""
        meta = f"{s.width}x{s.height}"
        if s.bracket >= 0:
            meta += f" · bracket {s.bracket} ({s.bracket_role} of {s.bracket_size})"
        if s.iso:
            meta += f" · ISO {s.iso}"
        if s.focal_mm:
            meta += f" · {s.focal_mm:.0f}mm"
        return f"""<figure class="card {sel}" data-cluster="{s.cluster}">
  <div class="thumbwrap">{'<img loading="lazy" src="' + t + '">' if t else '<div class=noimg>no preview</div>'}
    <span class="score">{s.score:.0f}</span>
    <span class="cl">grp {s.cluster}{'' if s.bracket < 0 else ' · brk ' + str(s.bracket)}</span>
    {'<span class="tick">PICK</span>' if sel else ''}
  </div>
  <figcaption>
    <div class="fn" title="{s.filename}">{s.filename}</div>
    <div class="meta">{meta}{' · ' + str(len(s.faces)) + ' face(s)' if s.faces else ''}</div>
    {why}
    <details><summary>metrics</summary>{bars}</details>
  </figcaption>
</figure>"""

    cards = "\n".join(card(s) for s in ordered)
    return_path = outdir / "report.html"
    return_path.write_text(f"""<!doctype html><meta charset="utf-8">
<title>Shortlist — {outdir.name}</title>
<style>
:root{{--bg:#0e0f12;--fg:#e8e9ec;--dim:#8b8f98;--line:#22242b;--go:#3ddc84;--warn:#ffb020}}
*{{box-sizing:border-box}}
body{{margin:0;background:var(--bg);color:var(--fg);font:14px/1.45 -apple-system,BlinkMacSystemFont,"Segoe UI",system-ui,sans-serif}}
header{{position:sticky;top:0;z-index:9;background:rgba(14,15,18,.92);backdrop-filter:blur(8px);border-bottom:1px solid var(--line);padding:18px 24px}}
h1{{margin:0 0 6px;font-size:18px;letter-spacing:-.01em}}
.sub{{color:var(--dim);font-size:13px}}
.controls{{margin-top:12px;display:flex;gap:8px;flex-wrap:wrap}}
button{{background:#191b21;color:var(--fg);border:1px solid var(--line);padding:6px 12px;border-radius:6px;cursor:pointer;font-size:13px}}
button.on{{background:var(--go);color:#08130c;border-color:var(--go);font-weight:600}}
main{{display:grid;grid-template-columns:repeat(auto-fill,minmax(230px,1fr));gap:16px;padding:20px 24px 60px}}
.card{{margin:0;background:#14161b;border:1px solid var(--line);border-radius:10px;overflow:hidden}}
.card.sel{{border-color:var(--go);box-shadow:0 0 0 1px var(--go) inset}}
.thumbwrap{{position:relative;aspect-ratio:3/2;background:#0a0b0e;display:flex;align-items:center;justify-content:center}}
.thumbwrap img{{width:100%;height:100%;object-fit:cover;display:block}}
.noimg{{color:var(--dim);font-size:12px}}
.score{{position:absolute;left:8px;top:8px;background:rgba(0,0,0,.75);padding:3px 8px;border-radius:20px;font-weight:700;font-size:13px}}
.cl{{position:absolute;right:8px;top:8px;background:rgba(0,0,0,.6);color:var(--dim);padding:2px 7px;border-radius:20px;font-size:11px}}
.tick{{position:absolute;right:8px;bottom:8px;background:var(--go);color:#08130c;font-weight:700;font-size:11px;padding:3px 8px;border-radius:4px;letter-spacing:.04em}}
figcaption{{padding:10px 12px 12px}}
.fn{{font-size:12px;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}}
.meta{{color:var(--dim);font-size:11px;margin-top:3px}}
.flags{{margin-top:6px;display:flex;gap:4px;flex-wrap:wrap}}
.flags em{{font-style:normal;background:#2a2113;color:var(--warn);font-size:10px;padding:2px 6px;border-radius:3px}}
details{{margin-top:8px}}
summary{{cursor:pointer;color:var(--dim);font-size:11px}}
.bar{{display:grid;grid-template-columns:96px 1fr 34px;align-items:center;gap:6px;margin-top:4px;font-size:10px;color:var(--dim)}}
.bar i{{display:block;height:4px;background:var(--go);border-radius:2px;opacity:.8}}
.bar b{{text-align:right;font-weight:500}}
.bar.na{{opacity:.35}}
.hide{{display:none}}
</style>
<header>
  <h1>Shortlist — {n_ok} images → {len(picked)} picks</h1>
  <div class="sub">{n_clusters} distinct moments · mode <b>{args.mode}</b> · genre <b>{args.genre}</b> · {Path(args.folder).name}</div>
  <div class="controls">
    <button id="ball" class="on">All ({n_ok})</button>
    <button id="bsel">Picks only ({len(picked)})</button>
    <button id="brej">Rejected ({n_ok - len(picked)})</button>
  </div>
</header>
<main id="grid">
{cards}
</main>
<script>
const g=document.getElementById('grid'),bs=[...document.querySelectorAll('.controls button')];
function show(f){{[...g.children].forEach(c=>{{const s=c.classList.contains('sel');
 c.classList.toggle('hide', f==='sel'? !s : f==='rej'? s : false);}});}}
bs.forEach(b=>b.onclick=()=>{{bs.forEach(x=>x.classList.remove('on'));b.classList.add('on');
 show(b.id==='bsel'?'sel':b.id==='brej'?'rej':'all');}});
</script>
""", encoding="utf-8")
    return return_path


def write_csv(shots: list[Shot], out: Path) -> None:
    cols = ["filename", "score", "cluster", "bracket", "bracket_role", "selected",
            "reasons", "keep_reasons", "n_faces",
            "width", "height", "iso", "focal_mm"] + [f"m_{k}" for k in METRIC_KEYS] + \
           ["raw_sharp_global", "raw_face_sharp", "raw_noise_sigma", "raw_tilt_deg", "error"]
    with out.open("w", newline="", encoding="utf-8") as fh:
        wr = csv.writer(fh)
        wr.writerow(cols)
        for s in sorted(shots, key=lambda s: -s.score):
            wr.writerow([
                s.filename, s.score, s.cluster,
                ("" if s.bracket < 0 else s.bracket), s.bracket_role,
                int(s.selected), "; ".join(s.reasons), "; ".join(s.keep_reasons),
                len(s.faces), s.width, s.height, s.iso, s.focal_mm,
                *[("" if s.m.get(k) is None else s.m[k]) for k in METRIC_KEYS],
                round(s.raw.get("sharp_global", 0), 2), round(s.raw.get("face_sharp", -1), 3),
                round(s.raw.get("noise_sigma", 0), 3), round(s.raw.get("tilt_deg", 0), 2), s.error,
            ])


# --------------------------------------------------------------------------
# main
# --------------------------------------------------------------------------

def find_images(folder: Path, recursive: bool) -> list[str]:
    it = folder.rglob("*") if recursive else folder.glob("*")
    exts = IMAGE_EXTS | RAW_EXTS
    return sorted(str(p) for p in it if p.is_file() and p.suffix.lower() in exts)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Cherry engine one - per-frame judgment and selection (prototype)")
    ap.add_argument("folder")
    ap.add_argument("--mode", choices=["top", "unique", "social"], default="top")
    ap.add_argument("--n", type=int, default=50, help="shortlist size for --mode top/social")
    ap.add_argument("--genre", choices=sorted(GENRE_PROFILES), default="general")
    ap.add_argument("--per-cluster", type=int, default=1,
                    help="keepers per near-duplicate group in 'unique' mode")
    ap.add_argument("--floor", type=float, default=40.0,
                    help="minimum score to be selectable (0-100); returns fewer "
                         "than --n rather than padding with bad frames")
    ap.add_argument("--out", default=None, help="output dir (default: <folder>/_shortlist)")
    ap.add_argument("--copy", action="store_true", help="copy selected files into out/picks")
    ap.add_argument("--recursive", action="store_true")
    ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 1))
    ap.add_argument("--no-thumbs", action="store_true")
    ap.add_argument("--no-brackets", action="store_true",
                    help="do not detect exposure brackets; treat every ladder "
                         "as an ordinary burst (the pre-v0.2 behaviour)")
    ap.add_argument("--limit", type=int, default=0, help="only process the first N files")
    ap.add_argument("--scorer-plugin", default=None,
                    help="python module path exposing score_images(paths)->{path: 0-1}")
    ap.add_argument("--open", action="store_true", help="open the HTML report when done")
    args = ap.parse_args(argv)

    folder = Path(args.folder).expanduser().resolve()
    if not folder.is_dir():
        print(f"error: {folder} is not a directory", file=sys.stderr)
        return 2

    paths = find_images(folder, args.recursive)
    if args.limit:
        paths = paths[: args.limit]
    if not paths:
        print("error: no images found", file=sys.stderr)
        return 2

    outdir = Path(args.out).expanduser().resolve() if args.out else folder / "_shortlist"
    outdir.mkdir(parents=True, exist_ok=True)
    thumbdir = outdir / "thumbs"
    if not args.no_thumbs:
        thumbdir.mkdir(exist_ok=True)

    genre = "social" if args.mode == "social" and args.genre == "general" else args.genre

    t0 = time.time()
    print(f"scanning {len(paths)} images  ·  genre={genre}  mode={args.mode}", flush=True)

    shots: list[Shot] = []
    if args.workers > 1 and len(paths) > 4:
        with ProcessPoolExecutor(max_workers=args.workers) as ex:
            futs = {ex.submit(analyse, p): p for p in paths}
            for i, fu in enumerate(as_completed(futs), 1):
                try:
                    shots.append(fu.result())
                except Exception as e:
                    shots.append(Shot(path=futs[fu], filename=Path(futs[fu]).name,
                                      ok=False, error=str(e)[:120]))
                if i % 25 == 0 or i == len(paths):
                    print(f"  analysed {i}/{len(paths)}", flush=True)
    else:
        for i, p in enumerate(paths, 1):
            shots.append(analyse(p))
            if i % 25 == 0 or i == len(paths):
                print(f"  analysed {i}/{len(paths)}", flush=True)

    shots.sort(key=lambda s: s.filename)
    n_bad = sum(1 for s in shots if not s.ok)

    n_clusters = cluster_shots(shots)
    print(f"grouped into {n_clusters} distinct moments", flush=True)

    n_brackets = 0 if args.no_brackets else detect_brackets(shots)
    if n_brackets:
        n_frames = sum(1 for s in shots if s.bracket >= 0)
        print(f"found {n_brackets} exposure bracket(s) covering {n_frames} frames "
              f"- each counts as one image, and the whole ladder ships together",
              flush=True)

    plugin_scores = None
    if args.scorer_plugin:
        try:
            import importlib
            mod = importlib.import_module(args.scorer_plugin)
            plugin_scores = mod.score_images([s.path for s in shots if s.ok])
            print(f"plugin '{args.scorer_plugin}' scored {len(plugin_scores)} images", flush=True)
        except Exception as e:
            print(f"warning: scorer plugin failed ({e}); continuing without it", file=sys.stderr)

    for s in shots:
        if s.ok:
            s.score = score_shot(s, genre, plugin_scores)

    picked = select(shots, args.mode, args.n, args.per_cluster, args.floor)
    for s in picked:
        s.selected = True
    explain_keeps(shots, genre)

    thumbs: dict[str, str] = {}
    if not args.no_thumbs:
        for s in shots:
            if not s.ok:
                continue
            name = f"{abs(hash(s.path)) % (10 ** 12):012d}.jpg"
            if write_thumb(s.path, thumbdir / name):
                thumbs[s.path] = f"thumbs/{name}"

    if args.copy:
        pickdir = outdir / "picks"
        pickdir.mkdir(exist_ok=True)
        for rank, s in enumerate(picked, 1):
            shutil.copy2(s.path, pickdir / f"{rank:03d}_{s.filename}")

    write_csv(shots, outdir / "scores.csv")
    (outdir / "result.json").write_text(json.dumps({
        "folder": str(folder), "genre": genre, "mode": args.mode,
        "analysed": len(shots), "failed": n_bad, "clusters": n_clusters,
        "selected": [s.filename for s in picked],
        "shots": [{**asdict(s), "hist": None} for s in shots],
    }, indent=1, default=str), encoding="utf-8")
    report = write_html(shots, picked, outdir, thumbs, args)

    dt = time.time() - t0
    rate = len(paths) / dt if dt else 0
    print(f"\ndone in {dt:.1f}s ({rate:.1f} img/s{', ' + str(n_bad) + ' failed' if n_bad else ''})")
    print(f"  picks   {len(picked)} / {len(shots)}")
    print(f"  report  {report}")
    print(f"  csv     {outdir / 'scores.csv'}")

    if args.open:
        try:
            import webbrowser
            webbrowser.open(report.as_uri())
        except Exception:
            pass
    return 0


# --------------------------------------------------------------------------
# PLUGIN CONTRACT
# --------------------------------------------------------------------------
# Drop in a learned scorer without touching this file. Create scorer_clip.py:
#
#   def score_images(paths: list[str]) -> dict[str, float]:
#       """Return {path: score in 0..1}. Batch however you like."""
#       ...
#
# then run:  python shortlist.py photos/ --scorer-plugin scorer_clip
#
# The plugin score is blended 40/60 with the classical score, so a weak or
# partial plugin degrades gracefully instead of taking the whole system down.
# Reasonable plugins: LAION aesthetic head on CLIP ViT-L/14, pyiqa's
# `clipiqa+` / `topiq_nr`, or a VLM called on cluster winners only.
# --------------------------------------------------------------------------

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