"""
Empty Shelf / Out-of-Stock (OOS) Slot Detector
Algorithm: Planogram Geometric Row-Clustering & Gap Analysis
Methodology:
1. Clusters YOLO product bounding boxes into horizontal shelf tiers/rows.
2. Filters out vertical architectural dividers / central pillars.
3. Identifies horizontal gaps exceeding expected product pitch.
4. Computes precise red bounding boxes for empty shelf slots.
"""

import cv2
import numpy as np
from typing import List, Dict, Any, Tuple, Optional


class EmptyShelfDetector:
    def __init__(
        self,
        min_gap_ratio: float = 0.75,
        min_row_products: int = 4,
        row_height_tol: float = 0.35,
        pillar_margin: int = 25
    ):
        """
        Args:
            min_gap_ratio: Minimum ratio of gap width to median item width in row (default 0.75).
            min_row_products: Minimum products required to establish a valid shelf row.
            row_height_tol: Vertical clustering tolerance relative to item height.
            pillar_margin: Buffer zone around detected vertical pillars/dividers.
        """
        self.min_gap_ratio = min_gap_ratio
        self.min_row_products = min_row_products
        self.row_height_tol = row_height_tol
        self.pillar_margin = pillar_margin

    def cluster_rows(self, boxes: List[List[int]]) -> List[List[List[int]]]:
        """Clusters bounding boxes [x1, y1, x2, y2] into horizontal shelf rows."""
        if not boxes:
            return []

        # Sort by vertical center
        sorted_boxes = sorted(boxes, key=lambda b: (b[1] + b[3]) / 2)
        rows: List[List[List[int]]] = []

        for b in sorted_boxes:
            yc = (b[1] + b[3]) / 2
            placed = False
            for r in rows:
                ryc = float(np.mean([(x[1] + x[3]) / 2 for x in r]))
                rh = float(np.mean([x[3] - x[1] for x in r]))
                if abs(yc - ryc) < self.row_height_tol * rh:
                    r.append(b)
                    placed = True
                    break
            if not placed:
                rows.append([b])

        # Filter rows with sufficient products and sort each row left-to-right
        valid_rows = []
        for r in rows:
            if len(r) >= self.min_row_products:
                r.sort(key=lambda b: b[0])
                valid_rows.append(r)

        return valid_rows

    def detect_vertical_pillars(
        self,
        rows: List[List[List[int]]],
        image_width: int
    ) -> List[Tuple[int, int]]:
        """
        Identifies structural vertical shelf dividers (pillars) across rows.
        A pillar exhibits consistent gaps across multiple vertical shelf rows.
        """
        if len(rows) < 2:
            return []

        # Count gap intersections across rows in 10px bins
        bin_size = 10
        num_bins = int(np.ceil(image_width / bin_size))
        gap_votes = np.zeros(num_bins, dtype=int)

        for row in rows:
            for j in range(len(row) - 1):
                x1 = row[j][2]
                x2 = row[j + 1][0]
                gap_w = x2 - x1
                # Architectural dividers typically span 40px - 180px
                if 40 <= gap_w <= 180:
                    start_bin = int(x1 // bin_size)
                    end_bin = min(num_bins, int(np.ceil(x2 / bin_size)))
                    gap_votes[start_bin:end_bin] += 1

        # A pillar should have gap votes in at least 50% of the shelf rows
        min_votes = max(2, int(len(rows) * 0.5))
        pillars = []
        in_pillar = False
        p_start = 0

        for b in range(num_bins):
            if gap_votes[b] >= min_votes:
                if not in_pillar:
                    in_pillar = True
                    p_start = b * bin_size
            else:
                if in_pillar:
                    in_pillar = False
                    p_end = b * bin_size
                    pillars.append((max(0, p_start - self.pillar_margin), min(image_width, p_end + self.pillar_margin)))

        if in_pillar:
            pillars.append((max(0, p_start - self.pillar_margin), image_width))

        return pillars

    def detect_empty_slots(
        self,
        image_shape: Tuple[int, int],
        product_boxes: List[Any],
        custom_pillars: Optional[List[Tuple[int, int]]] = None
    ) -> List[Dict[str, Any]]:
        """
        Detects empty shelf slots from product boxes.

        Args:
            image_shape: (height, width) of the shelf image.
            product_boxes: List of bounding boxes [x1, y1, x2, y2] or dicts with 'box' key.
            custom_pillars: Optional pre-defined pillar intervals [(x1, x2), ...].

        Returns:
            List of empty slot dictionaries:
            [
                {
                    "id": "empty_1",
                    "box": [x1, y1, x2, y2],
                    "width": w,
                    "height": h,
                    "missing_items_est": int,
                    "center": [cx, cy],
                    "row_index": int
                }, ...
            ]
        """
        img_h, img_w = image_shape[:2]
        if not product_boxes:
            return []

        # Standardize box format
        boxes: List[List[int]] = []
        for b in product_boxes:
            if isinstance(b, dict) and "box" in b:
                boxes.append([int(v) for v in b["box"]])
            elif isinstance(b, (list, tuple)) and len(b) >= 4:
                boxes.append([int(b[0]), int(b[1]), int(b[2]), int(b[3])])

        rows = self.cluster_rows(boxes)
        if not rows:
            return []

        # Detect structural vertical pillars/frames
        pillars = custom_pillars if custom_pillars is not None else self.detect_vertical_pillars(rows, img_w)

        empty_slots: List[Dict[str, Any]] = []
        slot_counter = 1

        for row_idx, row in enumerate(rows):
            med_y1 = int(np.median([b[1] for b in row]))
            med_y2 = int(np.median([b[3] for b in row]))
            med_w = float(np.median([b[2] - b[0] for b in row]))
            med_h = med_y2 - med_y1

            for j in range(len(row) - 1):
                x1 = row[j][2]
                x2 = row[j + 1][0]
                gap_w = x2 - x1

                if gap_w < self.min_gap_ratio * med_w:
                    continue

                # Check if this gap overlaps significantly with a known structural pillar
                is_pillar = False
                for px1, px2 in pillars:
                    overlap_start = max(x1, px1)
                    overlap_end = min(x2, px2)
                    if overlap_end > overlap_start:
                        overlap_len = overlap_end - overlap_start
                        if overlap_len >= 0.6 * gap_w or (x1 >= px1 and x2 <= px2) or (x1 <= px1 and x2 >= px2):
                            is_pillar = True
                            break

                if is_pillar:
                    continue

                missing_est = max(1, int(round(gap_w / med_w)))
                cx = int((x1 + x2) / 2)
                cy = int((med_y1 + med_y2) / 2)

                empty_slots.append({
                    "id": f"empty_{slot_counter}",
                    "box": [int(x1), int(med_y1), int(x2), int(med_y2)],
                    "width": int(gap_w),
                    "height": int(med_h),
                    "missing_items_est": missing_est,
                    "center": [cx, cy],
                    "row_index": row_idx + 1
                })
                slot_counter += 1

        return empty_slots

    def draw_empty_slots(
        self,
        image: np.ndarray,
        empty_slots: List[Dict[str, Any]],
        box_color: Tuple[int, int, int] = (0, 0, 255),
        thickness: int = 3,
        draw_fill: bool = True,
        draw_label: bool = False
    ) -> np.ndarray:
        """
        Draws red bounding boxes around empty shelf slots.

        Args:
            image: BGR numpy image.
            empty_slots: List of empty slot dictionaries.
            box_color: BGR color (default Red: (0, 0, 255)).
            thickness: Bounding box line thickness.
            draw_fill: Whether to render a semi-transparent fill tint.
            draw_label: Whether to render label text.
        """
        annotated = image.copy()

        if draw_fill and empty_slots:
            overlay = annotated.copy()
            for slot in empty_slots:
                x1, y1, x2, y2 = slot["box"]
                cv2.rectangle(overlay, (x1, y1), (x2, y2), box_color, -1)
            cv2.addWeighted(overlay, 0.22, annotated, 0.78, 0, annotated)

        for slot in empty_slots:
            x1, y1, x2, y2 = slot["box"]
            cv2.rectangle(annotated, (x1, y1), (x2, y2), box_color, thickness)
            if draw_label:
                est = slot.get("missing_items_est", 1)
                text = f"Empty (~{est})"
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = 0.5
                tw, th = cv2.getTextSize(text, font, font_scale, 1)[0]
                cv2.rectangle(annotated, (x1, max(0, y1 - th - 6)), (x1 + tw + 6, y1), box_color, -1)
                cv2.putText(annotated, text, (x1 + 3, y1 - 4), font, font_scale, (255, 255, 255), 1, cv2.LINE_AA)

        return annotated


def detect_and_draw_empty_slots(
    image_path: str,
    detector_weights: str = "runs/detect/train/weights/best.pt",
    output_path: Optional[str] = "detected_empty_slots.jpg"
) -> Tuple[np.ndarray, List[Dict[str, Any]]]:
    """
    Convenience function: runs YOLO detection, extracts empty shelf slots,
    and draws red bounding boxes.
    """
    from detector import ShelfProductDetector

    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"Failed to read image at: {image_path}")

    detector = ShelfProductDetector(detector_weights)
    det_res = detector.detect(image)
    boxes = det_res.get("boxes", [])

    empty_detector = EmptyShelfDetector()
    empty_slots = empty_detector.detect_empty_slots(image.shape, boxes)

    # Draw products in existing clean bounding boxes (green)
    annotated = detector.draw_detections(image, boxes, draw_labels=False)
    # Draw empty slots in RED
    annotated = empty_detector.draw_empty_slots(annotated, empty_slots, box_color=(0, 0, 255), thickness=3)

    if output_path:
        cv2.imwrite(output_path, annotated)

    return annotated, empty_slots