#!/usr/bin/env python3
"""Clean, merge, and regularise window detections from a segmentation model.

Windows are the most safety-critical keep-out for a painting robot, and the
semantic model gets them wrong in predictable ways: a window split by mullions
comes out as several fragments, thin frames erode or break, and diagonal or
perspective views never produce a clean rectangle. Raw connected components
inherit all of those flaws.

This module fixes the output geometry without retraining:

* ``merge_fragments`` joins pieces separated by up to a gap (mullion bars,
  occlusions) into one window;
* ``minimum_area_rectangle`` fits an oriented rectangle with rotating calipers,
  so a window stays a window at any viewing angle;
* ``window_instances`` runs clean -> merge -> fit -> filter and returns
  instances with both axis-aligned boxes and oriented quads;
* ``regularized_mask`` rasterises the fitted rectangles into a clean keep-out
  mask.

numpy + scipy only, so it is unit-testable and runs anywhere. See
``python3 window_postprocess.py --help`` for the CLI.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw
from scipy import ndimage


def disk(radius):
    size = int(radius) * 2 + 1
    ys, xs = np.mgrid[0:size, 0:size]
    return (xs - radius) ** 2 + (ys - radius) ** 2 <= radius ** 2


def clean_mask(mask, min_area=0, open_radius=0, close_radius=0):
    """Morphological tidy-up: close gaps, open specks, drop small components."""
    mask = np.asarray(mask, dtype=bool)
    if close_radius > 0:
        mask = ndimage.binary_closing(mask, structure=disk(close_radius))
    if open_radius > 0:
        mask = ndimage.binary_opening(mask, structure=disk(open_radius))
    if min_area > 0:
        labels, count = ndimage.label(mask, structure=np.ones((3, 3), dtype=int))
        if count:
            sizes = ndimage.sum(mask, labels, index=np.arange(1, count + 1))
            keep = np.flatnonzero(sizes >= min_area) + 1
            mask = np.isin(labels, keep)
    return mask


def merge_fragments(mask, gap):
    """Merge components separated by up to ``gap`` pixels (e.g. mullion bars)."""
    mask = np.asarray(mask, dtype=bool)
    if gap <= 0 or not mask.any():
        labels, _ = ndimage.label(mask, structure=np.ones((3, 3), dtype=int))
        return labels
    radius = max(1, int(round(gap / 2.0)))
    dilated = ndimage.binary_dilation(mask, structure=disk(radius))
    merged, _ = ndimage.label(dilated, structure=np.ones((3, 3), dtype=int))
    labels = np.where(mask, merged, 0)
    return labels


def convex_hull(points):
    """Monotone-chain hull of (N, 2) points, returned counter-clockwise."""
    unique = sorted({(float(x), float(y)) for x, y in points})
    if len(unique) <= 2:
        return np.asarray(unique, float)

    def cross(o, a, b):
        return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])

    lower = []
    for point in unique:
        while len(lower) >= 2 and cross(lower[-2], lower[-1], point) <= 0:
            lower.pop()
        lower.append(point)
    upper = []
    for point in reversed(unique):
        while len(upper) >= 2 and cross(upper[-2], upper[-1], point) <= 0:
            upper.pop()
        upper.append(point)
    return np.asarray(lower[:-1] + upper[:-1], float)


def minimum_area_rectangle(points):
    """Oriented rectangle of least area around points (rotating calipers).

    Returns ``(corners, width, height, angle_deg)`` where ``corners`` is a
    (4, 2) array. Robust to any orientation, unlike an axis-aligned box.
    """
    points = np.asarray(points, float)
    hull = convex_hull(points)
    if len(hull) < 3:
        x0, y0 = points.min(axis=0)
        x1, y1 = points.max(axis=0)
        corners = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], float)
        return corners, float(x1 - x0), float(y1 - y0), 0.0
    best = None
    count = len(hull)
    for i in range(count):
        p1, p2 = hull[i], hull[(i + 1) % count]
        edge = p2 - p1
        length = float(np.hypot(*edge))
        if length == 0:
            continue
        axis_x = edge / length
        axis_y = np.array([-axis_x[1], axis_x[0]])
        relative = hull - p1
        proj_x = relative @ axis_x
        proj_y = relative @ axis_y
        width = float(proj_x.max() - proj_x.min())
        height = float(proj_y.max() - proj_y.min())
        area = width * height
        if best is None or area < best[0]:
            center = p1 + (0.5 * (proj_x.min() + proj_x.max())) * axis_x \
                + (0.5 * (proj_y.min() + proj_y.max())) * axis_y
            half_w, half_h = width / 2.0, height / 2.0
            corners = np.array([center - half_w * axis_x - half_h * axis_y,
                                center + half_w * axis_x - half_h * axis_y,
                                center + half_w * axis_x + half_h * axis_y,
                                center - half_w * axis_x + half_h * axis_y])
            best = (area, corners, width, height, float(np.degrees(np.arctan2(axis_x[1], axis_x[0]))))
    return best[1], best[2], best[3], best[4]


def window_instances(mask, confidence=None, min_area=64, merge_gap=8, open_radius=1, close_radius=2,
                     min_fill=0.5, max_aspect=8.0, max_instances=200):
    """Detect clean window instances from a binary mask (and optional probability).

    Pipeline: morphological clean -> fragment merge -> oriented rectangle fit ->
    filter by area, fill ratio (a window fills its own rectangle), and aspect
    ratio. Returns ``(instances, labels)`` where ``labels`` is a per-pixel map
    whose positive values join each kept instance.
    """
    mask = clean_mask(np.asarray(mask, dtype=bool), min_area=0,
                      open_radius=open_radius, close_radius=close_radius)
    labels = merge_fragments(mask, merge_gap)
    instances = []
    for label in range(1, int(labels.max()) + 1 if labels.size else 0):
        component = labels == label
        area = int(component.sum())
        if area < min_area:
            continue
        ys, xs = np.nonzero(component)
        corners, width, height, angle = minimum_area_rectangle(np.c_[xs, ys])
        if width <= 0 or height <= 0:
            continue
        fill = area / (width * height)
        aspect = max(width, height) / max(1e-6, min(width, height))
        if fill < min_fill or aspect > max_aspect:
            continue
        score = float(confidence[component].mean()) if confidence is not None else None
        instances.append({
            "label": label,
            "bbox_xyxy": [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())],
            "quad_xy": [[round(float(x), 2), round(float(y), 2)] for x, y in corners],
            "center_xy": [round(float(xs.mean()), 2), round(float(ys.mean()), 2)],
            "size_xy": [round(float(width), 2), round(float(height), 2)],
            "angle_deg": round(float(angle), 2),
            "fill_ratio": round(float(fill), 4),
            "area_pixels": area,
            "mean_confidence": round(score, 4) if score is not None else None,
        })
    instances.sort(key=lambda item: item["area_pixels"], reverse=True)
    instances = instances[:max_instances]
    kept = {item["label"] for item in instances}
    kept_labels = np.where(np.isin(labels, list(kept) or [0]), labels, 0)
    return instances, kept_labels


def regularized_mask(instances, shape, use_quad=True):
    """Rasterise fitted rectangles into a clean binary window mask."""
    canvas = Image.new("L", (shape[1], shape[0]), 0)
    draw = ImageDraw.Draw(canvas)
    for instance in instances:
        if use_quad:
            draw.polygon([tuple(point) for point in instance["quad_xy"]], fill=255)
        else:
            x0, y0, x1, y1 = instance["bbox_xyxy"]
            draw.rectangle([x0, y0, x1, y1], fill=255)
    return np.asarray(canvas) > 0


def window_metrics(prediction, target):
    """Pixel precision/recall/IoU plus instance counts for window evaluation."""
    prediction, target = np.asarray(prediction, bool), np.asarray(target, bool)
    true_positive = int((prediction & target).sum())
    false_positive = int((prediction & ~target).sum())
    false_negative = int((~prediction & target).sum())
    precision = true_positive / (true_positive + false_positive) if true_positive + false_positive else 1.0
    recall = true_positive / (true_positive + false_negative) if true_positive + false_negative else 1.0
    iou = true_positive / (true_positive + false_positive + false_negative) if (true_positive + false_positive + false_negative) else 1.0
    return {"precision": round(precision, 4), "recall": round(recall, 4), "iou": round(iou, 4),
            "predicted_pixels": int(prediction.sum()), "target_pixels": int(target.sum())}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("window", type=Path, help="window probability (grayscale) or binary mask PNG")
    parser.add_argument("--threshold", type=float, default=0.5, help="probability threshold for a prob map")
    parser.add_argument("--min-area", type=int, default=64)
    parser.add_argument("--merge-gap", type=int, default=8, help="merge fragments separated by up to this many pixels")
    parser.add_argument("--open-radius", type=int, default=1)
    parser.add_argument("--close-radius", type=int, default=2)
    parser.add_argument("--min-fill", type=float, default=0.5)
    parser.add_argument("--max-aspect", type=float, default=8.0)
    parser.add_argument("--max-instances", type=int, default=200)
    parser.add_argument("--out-mask", type=Path, help="write the regularized window mask")
    parser.add_argument("--overlay", type=Path, help="write an overlay of detected windows")
    parser.add_argument("--json", type=Path, help="write instance JSON")
    args = parser.parse_args()

    raw = np.asarray(Image.open(args.window).convert("L"))
    probability = raw.astype(np.float32) / 255.0 if raw.max() > 1 else raw.astype(np.float32)
    mask = probability >= args.threshold
    instances, labels = window_instances(mask, confidence=probability, min_area=args.min_area,
                                         merge_gap=args.merge_gap, open_radius=args.open_radius,
                                         close_radius=args.close_radius, min_fill=args.min_fill,
                                         max_aspect=args.max_aspect, max_instances=args.max_instances)
    result = {"windows": len(instances), "instances": instances,
              "raw_pixels": int(mask.sum()),
              "regularized_pixels": int(regularized_mask(instances, mask.shape).sum())}
    rendered = json.dumps(result, indent=2)
    print(rendered)
    if args.json:
        args.json.parent.mkdir(parents=True, exist_ok=True)
        args.json.write_text(rendered + "\n", encoding="utf-8")
    if args.out_mask:
        args.out_mask.parent.mkdir(parents=True, exist_ok=True)
        Image.fromarray(regularized_mask(instances, mask.shape).astype(np.uint8) * 255).save(args.out_mask)
    if args.overlay:
        canvas = np.stack([raw] * 3, axis=-1).astype(np.uint8)
        canvas[mask] = (canvas[mask] * 0.4 + np.array([0, 0, 160])).astype(np.uint8)
        image = Image.fromarray(canvas)
        draw = ImageDraw.Draw(image)
        for instance in instances:
            draw.polygon([tuple(point) for point in instance["quad_xy"]], outline=(0, 255, 255))
        args.overlay.parent.mkdir(parents=True, exist_ok=True)
        image.save(args.overlay)
        print(f"overlay: {args.overlay}")


if __name__ == "__main__":
    main()