src/geometry/wall_vectorizer.py

"""
wall_vectorizer.py
------------------
Converts YOLOv8 segmentation masks into clean 2D wall polygons.

Pipeline per mask:
  1. Binarize mask
  2. Morphological cleanup (close gaps, remove noise)
  3. Find contours
  4. Approximate contours to simplified polygons (Douglas-Peucker)
  5. Filter by area and aspect ratio
  6. Return list of WallPolygon objects

Usage:
    from src.geometry.wall_vectorizer import WallVectorizer

    vectorizer = WallVectorizer()
    walls = vectorizer.extract(segmentation_result, image_shape)
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional
import cv2
import numpy as np


# ── Data structures ───────────────────────────────────────────────────────────

@dataclass
class WallPolygon:
    """A single vectorized wall or room boundary polygon."""
    class_id:   int
    class_name: str
    points:     list[tuple[int, int]]    # pixel coordinates (x, y)
    area:       float                    # pixel area
    bbox:       tuple[int, int, int, int]  # (x, y, w, h)
    confidence: float = 1.0

    @property
    def is_wall(self) -> bool:
        return self.class_id in (0, 1)  # OuterWall, InnerWall

    @property
    def is_room(self) -> bool:
        return self.class_id in (6, 7, 8, 9, 10, 11, 12)

    @property
    def centroid(self) -> tuple[float, float]:
        if not self.points:
            return (0.0, 0.0)
        xs = [p[0] for p in self.points]
        ys = [p[1] for p in self.points]
        return (sum(xs) / len(xs), sum(ys) / len(ys))

    def to_numpy(self) -> np.ndarray:
        """Return points as (N, 2) numpy array."""
        return np.array(self.points, dtype=np.int32)


@dataclass
class VectorizationResult:
    """All vectorized elements from one floor plan."""
    walls:    list[WallPolygon] = field(default_factory=list)
    rooms:    list[WallPolygon] = field(default_factory=list)
    doors:    list[WallPolygon] = field(default_factory=list)
    windows:  list[WallPolygon] = field(default_factory=list)
    other:    list[WallPolygon] = field(default_factory=list)
    image_shape: tuple[int, int] = (0, 0)

    @property
    def all_polygons(self) -> list[WallPolygon]:
        return self.walls + self.rooms + self.doors + self.windows + self.other

    @property
    def summary(self) -> dict:
        return {
            "walls":   len(self.walls),
            "rooms":   len(self.rooms),
            "doors":   len(self.doors),
            "windows": len(self.windows),
            "other":   len(self.other),
            "total":   len(self.all_polygons),
        }


# ── Vectorizer ────────────────────────────────────────────────────────────────

class WallVectorizer:
    """
    Converts segmentation masks into clean 2D vector polygons.

    Args:
        epsilon_factor:   Douglas-Peucker approximation factor
                          (fraction of arc length). Lower = more detail.
        min_area:         Discard polygons smaller than this (px²).
        morph_kernel:     Kernel size for morphological cleanup.
        simplify_walls:   Extra simplification pass for wall polygons.
    """

    # Which class_ids map to which category (0-indexed, background excluded)
    WALL_IDS    = {0, 1}          # OuterWall, InnerWall
    DOOR_IDS    = {3}             # Door
    WINDOW_IDS  = {2}             # Window
    ROOM_IDS    = {6, 7, 8, 9, 10, 11, 12}  # room types

    CLASS_NAMES = [
        "OuterWall", "InnerWall", "Window", "Door", "Stairs",
        "Railing", "Kitchen", "LivingRoom", "Bedroom", "Bathroom",
        "Corridor", "Balcony", "Garage",
    ]

    def __init__(
        self,
        epsilon_factor: float = 0.008,
        min_area: int = 200,
        morph_kernel: int = 3,
        simplify_walls: bool = True,
    ):
        self.epsilon_factor = epsilon_factor
        self.min_area = min_area
        self.morph_kernel = morph_kernel
        self.simplify_walls = simplify_walls

    def extract(
        self,
        segmentation_result,
        image_shape: Optional[tuple] = None,
    ) -> VectorizationResult:
        """
        Extract vector polygons from a SegmentationResult (Phase 2 output).

        Args:
            segmentation_result: FloorPlanPredictor result object.
            image_shape:         (H, W) of the source image.

        Returns:
            VectorizationResult with categorized polygons.
        """
        if image_shape is None:
            image_shape = segmentation_result.image_shape

        h, w = image_shape[:2]
        result = VectorizationResult(image_shape=(h, w))

        for element in segmentation_result.elements:
            if element.mask is None:
                continue

            polygons = self._mask_to_polygons(
                mask=element.mask,
                class_id=element.class_id,
                class_name=element.class_name,
                confidence=element.confidence,
                is_wall=(element.class_id in self.WALL_IDS),
            )

            for poly in polygons:
                if poly.class_id in self.WALL_IDS:
                    result.walls.append(poly)
                elif poly.class_id in self.DOOR_IDS:
                    result.doors.append(poly)
                elif poly.class_id in self.WINDOW_IDS:
                    result.windows.append(poly)
                elif poly.class_id in self.ROOM_IDS:
                    result.rooms.append(poly)
                else:
                    result.other.append(poly)

        return result

    def extract_from_mask(
        self,
        mask: np.ndarray,
        class_id: int,
        class_name: str,
        confidence: float = 1.0,
    ) -> list[WallPolygon]:
        """
        Extract polygons directly from a binary mask array.
        Useful for testing without a full SegmentationResult.
        """
        return self._mask_to_polygons(
            mask=mask,
            class_id=class_id,
            class_name=class_name,
            confidence=confidence,
            is_wall=(class_id in self.WALL_IDS),
        )

    # ── Internal helpers ──────────────────────────────────────────────────────

    def _mask_to_polygons(
        self,
        mask: np.ndarray,
        class_id: int,
        class_name: str,
        confidence: float,
        is_wall: bool,
    ) -> list[WallPolygon]:
        """Convert a binary mask to a list of simplified polygons."""

        # Ensure binary uint8
        binary = (mask > 127).astype(np.uint8) * 255

        # Morphological cleanup
        k = self.morph_kernel
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k, k))
        binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

        # Find external contours
        contours, _ = cv2.findContours(
            binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        polygons = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < self.min_area:
                continue

            # Douglas-Peucker simplification
            epsilon = self.epsilon_factor * cv2.arcLength(contour, closed=True)

            # Walls get extra simplification to remove noise
            if is_wall and self.simplify_walls:
                epsilon *= 1.5

            approx = cv2.approxPolyDP(contour, epsilon, closed=True)

            # Need at least 3 points for a valid polygon
            if len(approx) < 3:
                continue

            points = [(int(pt[0][0]), int(pt[0][1])) for pt in approx]
            x, y, w, h = cv2.boundingRect(contour)

            polygons.append(WallPolygon(
                class_id=class_id,
                class_name=class_name,
                points=points,
                area=float(area),
                bbox=(x, y, w, h),
                confidence=confidence,
            ))

        # Sort by area descending (largest first)
        polygons.sort(key=lambda p: p.area, reverse=True)
        return polygons

    def draw(
        self,
        image: np.ndarray,
        result: VectorizationResult,
        draw_labels: bool = True,
    ) -> np.ndarray:
        """
        Draw vectorized polygons on an image for visualization.

        Returns annotated BGR image.
        """
        if len(image.shape) == 2:
            canvas = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        else:
            canvas = image.copy()

        colors = {
            "wall":    (50,  50,  200),
            "door":    (50,  200, 200),
            "window":  (200, 180,  50),
            "room":    (50,  180,  80),
            "other":   (150, 150, 150),
        }

        def draw_poly(polys, color, label_prefix=""):
            for poly in polys:
                pts = np.array(poly.points, dtype=np.int32)
                cv2.polylines(canvas, [pts], isClosed=True,
                              color=color, thickness=2)
                if draw_labels:
                    cx, cy = int(poly.centroid[0]), int(poly.centroid[1])
                    cv2.putText(canvas, poly.class_name,
                                (cx, cy), cv2.FONT_HERSHEY_SIMPLEX,
                                0.4, color, 1, cv2.LINE_AA)

        draw_poly(result.walls,   colors["wall"])
        draw_poly(result.doors,   colors["door"])
        draw_poly(result.windows, colors["window"])
        draw_poly(result.rooms,   colors["room"])
        draw_poly(result.other,   colors["other"])

        return canvas