"""
Instance Processor — DefectDiffu Edition

Stripped-down version: DefectDiffu does NOT use 16x16 patch mappings,
so all patch-based artifact logic has been removed.

Retained utilities:
  - bbox ↔ mask helpers (for verification cropping)
  - IoU calculation
  - Visualization helpers
"""

import random
import torch
import numpy as np
from PIL import Image
from typing import List, Dict, Tuple, Optional, Union
import matplotlib.pyplot as plt
import matplotlib.patches as patches


class InstanceProcessor:
    """Utility class for detection post-processing and mask operations."""

    @staticmethod
    def calculate_iou(box1: Union[List, np.ndarray], box2: Union[List, np.ndarray]) -> float:
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        if x2 <= x1 or y2 <= y1:
            return 0.0
        intersection = (x2 - x1) * (y2 - y1)
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        union = area1 + area2 - intersection
        return intersection / union if union > 0 else 0.0

    @staticmethod
    def mask_from_bbox(bbox: Tuple[int, int, int, int], img_shape: Tuple[int, ...]) -> np.ndarray:
        """Create a binary mask from a bounding box."""
        h, w = img_shape[:2]
        mask = np.zeros((h, w), dtype=np.uint8)
        x1, y1, x2, y2 = bbox
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)
        if x2 > x1 and y2 > y1:
            mask[y1:y2, x1:x2] = 1
        return mask

    @staticmethod
    def get_bbox_from_mask(mask: np.ndarray, margin: int = 0) -> Tuple[int, int, int, int]:
        """Compute tight bounding box from binary mask, with optional margin."""
        ys, xs = np.where(mask > 0)
        if len(ys) == 0:
            return (0, 0, 0, 0)
        y1, y2 = ys.min(), ys.max()
        x1, x2 = xs.min(), xs.max()
        h, w = mask.shape
        x1 = max(0, x1 - margin)
        y1 = max(0, y1 - margin)
        x2 = min(w, x2 + margin)
        y2 = min(h, y2 + margin)
        return (x1, y1, x2, y2)

    @staticmethod
    def visualize_generation_result(
        original_image: np.ndarray,
        generated_image: np.ndarray,
        defect_mask: np.ndarray,
        output_path: str,
        title: str = "DefectDiffu Generation Result"
    ):
        """Create a 3-panel visualization: original, generated, mask overlay."""
        fig, axes = plt.subplots(1, 3, figsize=(18, 6))

        axes[0].imshow(original_image)
        axes[0].set_title("Original (Planning Reference)")
        axes[0].axis("off")

        axes[1].imshow(generated_image)
        axes[1].set_title("Generated Defect Image")
        axes[1].axis("off")

        axes[2].imshow(generated_image)
        axes[2].imshow(defect_mask, alpha=0.5, cmap="Reds")
        axes[2].set_title("Defect Mask Overlay")
        axes[2].axis("off")

        fig.suptitle(title, fontsize=14)
        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"[Viz] Saved result visualization to {output_path}")

    @staticmethod
    def resize_to_square(image: np.ndarray, size: int = 512) -> np.ndarray:
        """Resize image to square (DefectDiffu expects 512x512)."""
        pil_img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
        pil_img = pil_img.resize((size, size), Image.LANCZOS)
        return np.array(pil_img)