#!/usr/bin/env python3
"""Evaluate a rendered character against a generated blueprint.

The blueprint comes from `Tools/comfyui_character_blueprint.py`; the render comes from
the review pipeline. This tool isolates the figure in both, then reports how far the
model has drifted from the intended character:

    python3 Tools/character_match.py \
        --render Screenshots/demo/review-demo-third-person-front.png \
        --blueprint Art/characters/blueprints/front.png \
        --out /tmp/gigline/character-match.json \
        --montage /tmp/gigline/character-match.png

Reported per pair:

    aspect_delta     |height/width difference| between the two figures, in ratio units
    fill_delta       |figure area / bbox area| - bulk, catches a thin or bloated build
    coverage_delta   |figure area / frame area| - how much of the frame the figure takes
    mean_delta_e     mean CIE76 colour distance over the figure
    palette_delta_e  mean distance between the five dominant colours of each figure

Why these and not a similarity score: an absolute "match percentage" between a
photoreal generated sheet and a low-detail game body would be a number with no
meaning. The deltas are separable, so a report says *what* drifted - build, framing
or colour - instead of blending all three into one figure.

The figure mask is a heuristic and the tool says so in its report: for the blueprint
it is "differs from the corner colour", for the render it is "more saturated or darker
than the frame median", largest connected component. Pair the tool with a human look
at the montage; it narrows where to look, it does not replace the review.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys

from PIL import Image


# ---------------------------------------------------------------- colour maths

def srgb_to_linear(value: float) -> float:
    value /= 255.0
    return value / 12.92 if value <= 0.04045 else ((value + 0.055) / 1.055) ** 2.4


def rgb_to_lab(rgb) -> tuple[float, float, float]:
    """sRGB (0-255) to CIELAB under D65."""
    r, g, b = (srgb_to_linear(v) for v in rgb[:3])
    x = r * 0.4124 + g * 0.3576 + b * 0.1805
    y = r * 0.2126 + g * 0.7152 + b * 0.0722
    z = r * 0.0193 + g * 0.1192 + b * 0.9505
    # D65 white point
    xn, yn, zn = 0.95047, 1.0, 1.08883

    def f(t: float) -> float:
        return t ** (1.0 / 3.0) if t > 0.008856 else (7.787 * t) + (16.0 / 116.0)

    fx, fy, fz = f(x / xn), f(y / yn), f(z / zn)
    return (116.0 * fy - 16.0, 500.0 * (fx - fy), 200.0 * (fy - fz))


def delta_e76(lab_a, lab_b) -> float:
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(lab_a, lab_b)))


# ------------------------------------------------------------------- masking

def luminance(rgb) -> float:
    return 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2]


def _subject_mask(image: Image.Image, saturation_threshold: int = 20,
                  dark_ratio: float = 0.62) -> Image.Image:
    """Mask the figure as "more saturated than the backdrop, or notably darker".

    Both inputs have a desaturated backdrop - grey studio sweep in the blueprint,
    beige pavement and pale walls in the render - and a figure that carries colour
    (skin, denim, hair). Saturation separates them without needing a clean plate.
    The previous version masked the blueprint by "differs from the corner colour",
    which swallowed a gradient backdrop whole.
    """
    rgb = image.convert("RGB")
    pixels = rgb.load()
    lums = [luminance(pixels[x, y])
            for y in range(0, rgb.height, 4) for x in range(0, rgb.width, 4)]
    lums.sort()
    median = lums[len(lums) // 2] if lums else 128.0

    mask = Image.new("L", rgb.size, 0)
    out = mask.load()
    for y in range(rgb.height):
        for x in range(rgb.width):
            r, g, b = pixels[x, y]
            saturation = max(r, g, b) - min(r, g, b)
            if saturation > saturation_threshold or luminance((r, g, b)) < median * dark_ratio:
                out[x, y] = 255
    return mask


def figure_mask_from_background(image: Image.Image, tolerance: int = 28):
    """Figure mask for a blueprint sheet: a coloured subject on a grey sweep."""
    return _largest_component(_subject_mask(image, saturation_threshold=18,
                                            dark_ratio=0.55))


def figure_mask_from_frame(image: Image.Image):
    """Figure mask for a rendered frame: the body against pavement and stucco."""
    return _largest_component(_subject_mask(image, saturation_threshold=26,
                                            dark_ratio=0.72))


def _largest_component(mask: Image.Image) -> Image.Image:
    """Keep the largest plausible figure blob.

    "Largest" alone is wrong twice over: in both inputs the backdrop is one huge
    connected region, so a plain largest-blob pick returned the background and the
    first report measured the whole frame. The filter is size-based instead - a figure
    occupies between 0.2 % and 55 % of the frame - which rejects the backdrop on one
    side and road markings on the other.
    """
    width, height = mask.size
    frame_area = width * height
    source = mask.load()
    seen = bytearray(width * height)
    best = []      # largest blob fully inside the frame
    usable = []    # largest plausible blob, border or not
    fallback = []  # largest blob at all, used when nothing is plausible
    best_score = 0

    for start_y in range(0, height, 3):
        for start_x in range(0, width, 3):
            index = start_y * width + start_x
            if seen[index] or source[start_x, start_y] == 0:
                continue
            stack = [(start_x, start_y)]
            blob = []
            seen[index] = 1
            while stack:
                x, y = stack.pop()
                blob.append((x, y))
                for nx, ny in ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)):
                    if 0 <= nx < width and 0 <= ny < height:
                        nindex = ny * width + nx
                        if not seen[nindex] and source[nx, ny] != 0:
                            seen[nindex] = 1
                            stack.append((nx, ny))
            if len(blob) > len(fallback):
                fallback = blob
            share = len(blob) / frame_area
            if not (0.002 < share < 0.55):
                continue
            # Blobs glued to the frame edge are backdrop: the sky band in a render and
            # the studio sweep in a sheet both touch the border and both beat the
            # figure on raw area. Prefer fully enclosed blobs, fall back by area.
            xs = [p[0] for p in blob]
            ys = [p[1] for p in blob]
            enclosed = (min(xs) > 1 and min(ys) > 1
                        and max(xs) < width - 2 and max(ys) < height - 2)
            if enclosed and len(blob) > best_score:
                best_score = len(blob)
                best = blob
            elif len(blob) > len(usable):
                usable = blob

    keep = Image.new("L", mask.size, 0)
    out = keep.load()
    for x, y in (best or usable or fallback):
        out[x, y] = 255
    return keep


def bbox(mask: Image.Image):
    return mask.getbbox()


def mask_stats(mask: Image.Image) -> dict:
    box = mask.getbbox()
    if not box:
        return {"area": 0, "bbox": None, "aspect": None, "fill": None, "coverage": None}
    x0, y0, x1, y1 = box
    histogram = mask.histogram()
    area = sum(histogram[1:])
    box_area = max(1, (x1 - x0) * (y1 - y0))
    return {
        "area": area,
        "bbox": [x0, y0, x1, y1],
        "aspect": (y1 - y0) / max(1, (x1 - x0)),
        "fill": area / box_area,
        "coverage": area / (mask.width * mask.height),
    }


def masked_colours(image: Image.Image, mask: Image.Image, limit=20000):
    rgb = image.convert("RGB")
    pixels = rgb.load()
    values = mask.load()
    collected = []
    step = max(1, (rgb.width * rgb.height) // limit)
    index = 0
    for y in range(rgb.height):
        for x in range(rgb.width):
            index += 1
            if index % step:
                continue
            if values[x, y]:
                collected.append(pixels[x, y])
    return collected


def dominant(counts: dict, count: int = 5):
    """Coarse palette: quantise to 32 levels per channel, then take the top bins."""
    bins: dict[tuple, int] = {}
    for colour in counts:
        key = (colour[0] // 32, colour[1] // 32, colour[2] // 32)
        bins[key] = bins.get(key, 0) + 1
    ranked = sorted(bins.items(), key=lambda item: -item[1])[:count]
    return [((k[0] * 32) + 16, (k[1] * 32) + 16, (k[2] * 32) + 16) for k, _ in ranked]


def palette_delta_e(palette_a, palette_b) -> float:
    if not palette_a or not palette_b:
        return float("nan")
    labs_b = [rgb_to_lab(c) for c in palette_b]
    total = 0.0
    for colour in palette_a:
        lab = rgb_to_lab(colour)
        total += min(delta_e76(lab, other) for other in labs_b)
    return total / len(palette_a)


# -------------------------------------------------------------------- report

def evaluate(render_path: str, blueprint_path: str, montage_path: str | None = None,
             render_box: tuple[float, float, float, float] | None = None) -> dict:
    render_full = Image.open(render_path).convert("RGB")
    blueprint = Image.open(blueprint_path).convert("RGB")

    # The figure box is explicit by default. Auto-segmentation on these shots failed
    # in a way worth recording rather than papering over: the backdrop (sky band, pale
    # wall, dark pavement) is more saturated and far larger than the subject, so every
    # size-based rule picked a backdrop blob. A box read off the shot is honest and
    # cheap; a background plate rendered without the character would be the proper fix.
    width, height = render_full.size
    box = render_box or (0.15, 0.02, 0.85, 0.99)
    region = (int(width * box[0]), int(height * box[1]),
              int(width * box[2]), int(height * box[3]))
    render = render_full.crop(region)

    render_mask = figure_mask_from_frame(render)
    blueprint_mask = figure_mask_from_background(blueprint)

    render_stats = mask_stats(render_mask)
    blueprint_stats = mask_stats(blueprint_mask)

    render_colours = masked_colours(render, render_mask)
    blueprint_colours = masked_colours(blueprint, blueprint_mask)

    def mean_lab(colours):
        labs = [rgb_to_lab(c) for c in colours]
        if not labs:
            return None
        return tuple(sum(l[i] for l in labs) / len(labs) for i in range(3))

    lab_render = mean_lab(render_colours)
    lab_blueprint = mean_lab(blueprint_colours)
    mean_de = (delta_e76(lab_render, lab_blueprint)
               if lab_render and lab_blueprint else float("nan"))

    palette_render = dominant(render_colours)
    palette_blueprint = dominant(blueprint_colours)

    def delta(a, b):
        if a is None or b is None:
            return None
        return abs(a - b)

    report = {
        "render": os.path.abspath(render_path),
        "blueprint": os.path.abspath(blueprint_path),
        "render_mask": "saturated-or-dark vs frame median, largest component",
        "blueprint_mask": "differs from corner colour",
        "render_region": list(region),
        "render_figure": render_stats,
        "blueprint_figure": blueprint_stats,
        "aspect_delta": delta(render_stats["aspect"], blueprint_stats["aspect"]),
        "fill_delta": delta(render_stats["fill"], blueprint_stats["fill"]),
        "coverage_delta": delta(render_stats["coverage"], blueprint_stats["coverage"]),
        "mean_lab_render": lab_render,
        "mean_lab_blueprint": lab_blueprint,
        "mean_delta_e": mean_de,
        "palette_render": palette_render,
        "palette_blueprint": palette_blueprint,
        "palette_delta_e": palette_delta_e(palette_render, palette_blueprint),
    }

    if montage_path:
        report["montage"] = write_montage(render, render_mask, blueprint, blueprint_mask,
                                          montage_path)
    return report


def write_montage(render, render_mask, blueprint, blueprint_mask, path, height=640):
    """Render crop and blueprint crop side by side, same height, for a human look."""
    def crop(image, mask):
        box = mask.getbbox()
        return image.crop(box) if box else image

    left = crop(render, render_mask)
    right = crop(blueprint, blueprint_mask)
    left = left.resize((max(1, int(left.width * height / max(1, left.height))), height))
    right = right.resize((max(1, int(right.width * height / max(1, right.height))), height))
    sheet = Image.new("RGB", (left.width + right.width + 12, height), (20, 20, 20))
    sheet.paste(left, (0, 0))
    sheet.paste(right, (left.width + 12, 0))
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    sheet.save(path)
    return os.path.abspath(path)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--render", required=True)
    ap.add_argument("--blueprint", required=True)
    ap.add_argument("--out", default="/tmp/gigline/character-match.json")
    ap.add_argument("--montage", default=None)
    ap.add_argument("--render-box", default=None,
                    help="figure box as x0,y0,x1,y1 fractions of the frame (default: central 70 percent)")
    args = ap.parse_args(argv)

    for path in (args.render, args.blueprint):
        if not os.path.exists(path):
            print(f"error: no such image: {path}", file=sys.stderr)
            return 2

    box = None
    if args.render_box:
        box = tuple(float(v) for v in args.render_box.split(","))
        if len(box) != 4:
            print("error: --render-box needs four comma-separated fractions", file=sys.stderr)
            return 2
    report = evaluate(args.render, args.blueprint, args.montage, box)
    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
    with open(args.out, "w") as handle:
        json.dump(report, handle, indent=2)

    def show(value, digits=3):
        return "n/a" if value is None or (isinstance(value, float) and math.isnan(value)) \
            else f"{value:.{digits}f}"

    print(f"render   {report['render_figure']}")
    print(f"blueprint {report['blueprint_figure']}")
    print(f"aspect delta    {show(report['aspect_delta'])}")
    print(f"fill delta      {show(report['fill_delta'])}")
    print(f"coverage delta  {show(report['coverage_delta'])}")
    print(f"mean deltaE     {show(report['mean_delta_e'], 1)}")
    print(f"palette deltaE  {show(report['palette_delta_e'], 1)}")
    print(f"report -> {args.out}")
    return 0


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