# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np
import torchvision.transforms as T  # noqa: N812
from supervision import BoxAnnotator, Color, Detections, LabelAnnotator, MaskAnnotator
from torch import Tensor
from torch.utils.data import DataLoader

from rfdetr.utilities.box_ops import box_cxcywh_to_xyxy
from rfdetr.utilities.logger import get_logger
from rfdetr.visualize.training import _ensure_headless_backend

if TYPE_CHECKING:
    from matplotlib.axes import Axes

logger = get_logger()


class DatasetGridSaver:
    """Utility for saving 3x3 image grids to visualize augmentation effects.

    Args:
        data_loader: Dataloader of the dataset to sample images from.
        output_dir: Directory where grid images will be saved.
        max_batches: Number of batches to draw samples from.
        dataset_type: Dataset split label, e.g. ``'train'`` or ``'val'``.
    """

    def __init__(
        self, data_loader: DataLoader[Any], output_dir: Path, max_batches: int = 3, dataset_type: str = "train"
    ) -> None:
        self.data_loader = data_loader
        self.output_dir = output_dir
        self.max_batches = max_batches
        self.dataset_type = dataset_type
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def save_grid(self) -> None:
        """Create and save image grids to ``output_dir``.

        Each grid is a 3x3 JPEG containing up to 9 images from a single batch, with bounding boxes and class labels
        drawn on top. Instance masks are also drawn, before the boxes and labels, for any sample whose target includes a
        ``'masks'`` entry.
        """
        try:
            import matplotlib

            _ensure_headless_backend(matplotlib)
            import matplotlib.pyplot as plt
        except ImportError as exc:
            raise ImportError(
                "matplotlib is required for dataset grids. Install it with: pip install matplotlib"
            ) from exc

        inv_normalize = T.Normalize(
            mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
            std=[1 / 0.229, 1 / 0.224, 1 / 0.225],
        )
        mask_annotator = MaskAnnotator(opacity=0.45)
        box_annotator = BoxAnnotator(thickness=2)
        label_annotator = LabelAnnotator(
            text_color=Color.BLACK,
            text_scale=0.5,
            text_padding=3,
        )

        for batch_idx, (sample, target) in enumerate(self.data_loader):
            if batch_idx >= self.max_batches:
                break

            fig, axes = plt.subplots(3, 3, figsize=(12, 12))
            fig.suptitle(f"{self.dataset_type} dataset, batch {batch_idx}")
            axes = axes.flatten()

            num_plotted = 0
            for single_image, single_target in zip(sample.tensors, target):
                if num_plotted >= 9:
                    break
                self._annotate_and_plot(
                    single_image,
                    single_target,
                    axes[num_plotted],
                    inv_normalize,
                    mask_annotator,
                    box_annotator,
                    label_annotator,
                )
                num_plotted += 1

            for i in range(num_plotted, 9):
                axes[i].axis("off")

            fig.tight_layout()
            plt.savefig(self.output_dir / f"{self.dataset_type}_batch{batch_idx}_grid.jpg", dpi=200)
            plt.close()

        logger.info(f"Saved {self.dataset_type} grids with augmented images to: {self.output_dir.resolve()}")

    @staticmethod
    def _annotate_and_plot(
        single_image: Tensor,
        single_target: dict[str, Any],
        ax: Axes,
        inv_normalize: T.Normalize,
        mask_annotator: MaskAnnotator,
        box_annotator: BoxAnnotator,
        label_annotator: LabelAnnotator,
    ) -> None:
        """De-normalize a single image tensor, annotate it with masks, boxes, and labels, and plot it on ``ax``.

        Args:
            single_image: Normalized image tensor of shape ``(C, H, W)``.
            single_target: Target dict with keys ``'size'``, ``'boxes'`` (cx,cy,wh normalized), ``'labels'``, and
                optionally ``'masks'``.
            ax: Matplotlib axis to plot the annotated image on.
            inv_normalize: Inverse normalization transform to convert the tensor back to pixel values.
            mask_annotator: ``MaskAnnotator`` instance for drawing instance masks.
            box_annotator: ``BoxAnnotator`` instance for drawing bounding boxes.
            label_annotator: ``LabelAnnotator`` instance for drawing class labels.
        """
        from PIL import Image as PILImage

        resized_size = single_target["size"]
        if isinstance(resized_size, Tensor):
            resized_size = resized_size.detach().cpu()
        h, w = int(resized_size[0]), int(resized_size[1])

        de_normalized_img = inv_normalize(single_image[:, :h, :w])
        if isinstance(de_normalized_img, Tensor):
            de_normalized_img = de_normalized_img.detach().cpu().numpy()
        scene = PILImage.fromarray(
            np.ascontiguousarray((np.clip(de_normalized_img.transpose(1, 2, 0), 0.0, 1.0) * 255).astype(np.uint8))
        )

        if len(single_target["boxes"]) > 0:
            labels_tensor = single_target["labels"]
            if isinstance(labels_tensor, Tensor):
                class_ids = labels_tensor.detach().cpu().numpy().astype(int)
            else:
                class_ids = np.asarray(labels_tensor, dtype=int)

            boxes = single_target["boxes"]
            if isinstance(boxes, Tensor):
                boxes_iter = boxes.detach().cpu()
            else:
                boxes_iter = boxes

            xyxy = np.asarray(
                [[b[0] * w, b[1] * h, b[2] * w, b[3] * h] for box in boxes_iter for b in [box_cxcywh_to_xyxy(box)]],
                dtype=np.float32,
            )
            masks = DatasetGridSaver._extract_masks(single_target, image_shape=(h, w), num_instances=len(class_ids))
            detections = Detections(xyxy=xyxy, mask=masks, class_id=class_ids)
            labels = [str(c) for c in class_ids]
            if masks is not None:
                scene = mask_annotator.annotate(scene=scene, detections=detections)
            scene = box_annotator.annotate(scene=scene, detections=detections)
            scene = label_annotator.annotate(scene=scene, detections=detections, labels=labels)

        ax.imshow(scene)
        ax.axis("off")

    @staticmethod
    def _extract_masks(
        single_target: dict[str, Any],
        image_shape: tuple[int, int],
        num_instances: int,
    ) -> np.ndarray[Any, Any] | None:
        """Return target masks as ``(N, H, W)`` booleans aligned to the rendered image shape.

        Args:
            single_target: Target dict that may contain a ``'masks'`` entry — a tensor or array-like of shape
                ``(N, H', W')`` where ``H'``/``W'`` need not match ``image_shape``.
            image_shape: ``(height, width)`` of the rendered image the masks must align to.
            num_instances: Number of detected instances (boxes/labels) the returned masks must match.

        Returns:
            A ``(num_instances, height, width)`` boolean array aligned to ``image_shape``, or ``None`` when masks
            are absent, malformed, or insufficient for ``num_instances``.
        """
        masks = single_target.get("masks")
        if masks is None or (hasattr(masks, "ndim") and masks.ndim == 0):
            return None
        if len(masks) == 0:
            return None

        if isinstance(masks, Tensor):
            masks_np = masks.detach().cpu().numpy()
        else:
            masks_np = np.asarray(masks)

        if masks_np.ndim != 3:
            logger.warning("Skipping dataset-grid masks with invalid shape %s; expected (N, H, W).", masks_np.shape)
            return None

        if masks_np.shape[0] < num_instances:
            logger.warning(
                "Skipping dataset-grid masks with %d masks for %d instances.",
                masks_np.shape[0],
                num_instances,
            )
            return None

        masks_np = masks_np[:num_instances].astype(bool, copy=False)
        height, width = image_shape
        if masks_np.shape[1:] == image_shape:
            return masks_np

        aligned_masks = np.zeros((masks_np.shape[0], height, width), dtype=bool)
        copy_height = min(height, masks_np.shape[1])
        copy_width = min(width, masks_np.shape[2])
        aligned_masks[:, :copy_height, :copy_width] = masks_np[:, :copy_height, :copy_width]
        return aligned_masks
