# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# =========================================================================
# Adapted from https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py
# which has the following license...
# https://github.com/pytorch/vision/blob/main/LICENSE

# BSD 3-Clause License

# Copyright (c) Soumith Chintala 2016,
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:

# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.

# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.

# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
"""
This script is adapted from
https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py
"""

from __future__ import annotations

from collections.abc import Sequence

import torch
from torch import Tensor, nn

from monai.utils import ensure_tuple
from monai.utils.misc import issequenceiterable
from monai.utils.module import look_up_option


class AnchorGenerator(nn.Module):
    """
    This module is modified from torchvision to support both 2D and 3D images.

    Module that generates anchors for a set of feature maps and
    image sizes.

    The module support computing anchors at multiple sizes and aspect ratios
    per feature map.

    sizes and aspect_ratios should have the same number of elements, and it should
    correspond to the number of feature maps.

    sizes[i] and aspect_ratios[i] can have an arbitrary number of elements.
    For 2D images, anchor width and height w:h = 1:aspect_ratios[i,j]
    For 3D images, anchor width, height, and depth w:h:d = 1:aspect_ratios[i,j,0]:aspect_ratios[i,j,1]

    AnchorGenerator will output a set of sizes[i] * aspect_ratios[i] anchors
    per spatial location for feature map i.

    Args:
        sizes: base size of each anchor.
            len(sizes) is the number of feature maps, i.e., the number of output levels for
            the feature pyramid network (FPN).
            Each element of ``sizes`` is a Sequence which represents several anchor sizes for each feature map.
        aspect_ratios: the aspect ratios of anchors. ``len(aspect_ratios) = len(sizes)``.
            For 2D images, each element of ``aspect_ratios[i]`` is a Sequence of float.
            For 3D images, each element of ``aspect_ratios[i]`` is a Sequence of 2 value Sequence.
        indexing: choose from {``'ij'``, ``'xy'``}, optional,
            Matrix (``'ij'``, default and recommended) or Cartesian (``'xy'``) indexing of output.

            - Matrix (``'ij'``, default and recommended) indexing keeps the original axis not changed.
            - To use other monai detection components, please set ``indexing = 'ij'``.
            - Cartesian (``'xy'``) indexing swaps axis 0 and 1.
            - For 2D cases, monai ``AnchorGenerator(sizes, aspect_ratios, indexing='xy')`` and
              ``torchvision.models.detection.anchor_utils.AnchorGenerator(sizes, aspect_ratios)`` are equivalent.


    Reference:.
        https://github.com/pytorch/vision/blob/release/0.12/torchvision/models/detection/anchor_utils.py

    Example:
        .. code-block:: python

            # 2D example inputs for a 2-level feature maps
            sizes = ((10,12,14,16), (20,24,28,32))
            base_aspect_ratios = (1., 0.5,  2.)
            aspect_ratios = (base_aspect_ratios, base_aspect_ratios)
            anchor_generator = AnchorGenerator(sizes, aspect_ratios)

            # 3D example inputs for a 2-level feature maps
            sizes = ((10,12,14,16), (20,24,28,32))
            base_aspect_ratios = ((1., 1.), (1., 0.5), (0.5, 1.), (2., 2.))
            aspect_ratios = (base_aspect_ratios, base_aspect_ratios)
            anchor_generator = AnchorGenerator(sizes, aspect_ratios)
    """

    __annotations__ = {"cell_anchors": list[torch.Tensor]}

    def __init__(
        self,
        sizes: Sequence[Sequence[int]] = ((20, 30, 40),),
        aspect_ratios: Sequence = (((0.5, 1), (1, 0.5)),),
        indexing: str = "ij",
    ) -> None:
        super().__init__()

        if not issequenceiterable(sizes[0]):
            self.sizes = tuple((s,) for s in sizes)
        else:
            self.sizes = ensure_tuple(sizes)
        if not issequenceiterable(aspect_ratios[0]):
            aspect_ratios = (aspect_ratios,) * len(self.sizes)

        if len(self.sizes) != len(aspect_ratios):
            raise ValueError(
                "len(sizes) and len(aspect_ratios) should be equal. \
                It represents the number of feature maps."
            )

        spatial_dims = len(ensure_tuple(aspect_ratios[0][0])) + 1
        spatial_dims = look_up_option(spatial_dims, [2, 3])
        self.spatial_dims = spatial_dims

        self.indexing = look_up_option(indexing, ["ij", "xy"])

        self.aspect_ratios = aspect_ratios
        self.cell_anchors = [
            self.generate_anchors(size, aspect_ratio) for size, aspect_ratio in zip(self.sizes, aspect_ratios)
        ]

    # This comment comes from torchvision.
    # TODO: https://github.com/pytorch/pytorch/issues/26792
    # For every (aspect_ratios, scales) combination, output a zero-centered anchor with those values.
    # (scales, aspect_ratios) are usually an element of zip(self.scales, self.aspect_ratios)
    # This method assumes aspect ratio = height / width for an anchor.
    def generate_anchors(
        self,
        scales: Sequence,
        aspect_ratios: Sequence,
        dtype: torch.dtype = torch.float32,
        device: torch.device | None = None,
    ) -> torch.Tensor:
        """
        Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map.

        Args:
            scales: a sequence which represents several anchor sizes for the current feature map.
            aspect_ratios: a sequence which represents several aspect_ratios for the current feature map.
                For 2D images, it is a Sequence of float aspect_ratios[j],
                anchor width and height w:h = 1:aspect_ratios[j].
                For 3D images, it is a Sequence of 2 value Sequence aspect_ratios[j,0] and aspect_ratios[j,1],
                anchor width, height, and depth w:h:d = 1:aspect_ratios[j,0]:aspect_ratios[j,1]
            dtype: target data type of the output Tensor.
            device: target device to put the output Tensor data.

            Returns:
                For each s in scales, returns [s, s*aspect_ratios[j]] for 2D images,
                and [s, s*aspect_ratios[j,0],s*aspect_ratios[j,1]] for 3D images.
        """
        scales_t = torch.as_tensor(scales, dtype=dtype, device=device)  # sized (N,)
        aspect_ratios_t = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)  # sized (M,) or (M,2)
        if (self.spatial_dims >= 3) and (len(aspect_ratios_t.shape) != 2):
            raise ValueError(
                f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \
                {len(aspect_ratios_t.shape) - 1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}."
            )

        if (self.spatial_dims >= 3) and (aspect_ratios_t.shape[1] != self.spatial_dims - 1):
            raise ValueError(
                f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \
                shape (_,{self.spatial_dims - 1}). But got aspect_ratios with shape {aspect_ratios_t.shape}."
            )

        # if 2d, w:h = 1:aspect_ratios
        if self.spatial_dims == 2:
            area_scale = torch.sqrt(aspect_ratios_t)
            w_ratios = 1 / area_scale
            h_ratios = area_scale
        # if 3d, w:h:d = 1:aspect_ratios[:,0]:aspect_ratios[:,1]
        else:
            area_scale = torch.pow(aspect_ratios_t[:, 0] * aspect_ratios_t[:, 1], 1 / 3.0)
            w_ratios = 1 / area_scale
            h_ratios = aspect_ratios_t[:, 0] / area_scale
            d_ratios = aspect_ratios_t[:, 1] / area_scale

        ws = (w_ratios[:, None] * scales_t[None, :]).view(-1)
        hs = (h_ratios[:, None] * scales_t[None, :]).view(-1)
        if self.spatial_dims == 2:
            base_anchors = torch.stack([-ws, -hs, ws, hs], dim=1) / 2.0
        else:  # elif self.spatial_dims == 3:
            ds = (d_ratios[:, None] * scales_t[None, :]).view(-1)
            base_anchors = torch.stack([-ws, -hs, -ds, ws, hs, ds], dim=1) / 2.0

        return base_anchors.round()

    def set_cell_anchors(self, dtype: torch.dtype, device: torch.device) -> None:
        """
        Convert each element in self.cell_anchors to ``dtype`` and send to ``device``.
        """
        self.cell_anchors = [cell_anchor.to(dtype=dtype, device=device) for cell_anchor in self.cell_anchors]

    def num_anchors_per_location(self):
        """
        Return number of anchor shapes for each feature map.
        """
        return [c.shape[0] for c in self.cell_anchors]

    def grid_anchors(self, grid_sizes: list[list[int]], strides: list[list[Tensor]]) -> list[Tensor]:
        """
        Every combination of (a, (g, s), i) in (self.cell_anchors, zip(grid_sizes, strides), 0:spatial_dims)
        corresponds to a feature map.
        It outputs g[i] anchors that are s[i] distance apart in direction i, with the same dimensions as a.

        Args:
            grid_sizes: spatial size of the feature maps
            strides: strides of the feature maps regarding to the original image

        Example:
            .. code-block:: python

                grid_sizes = [[100,100],[50,50]]
                strides = [[torch.tensor(2),torch.tensor(2)], [torch.tensor(4),torch.tensor(4)]]
        """
        anchors = []
        cell_anchors = self.cell_anchors
        if cell_anchors is None:
            raise AssertionError

        if not (len(grid_sizes) == len(strides) == len(cell_anchors)):
            raise ValueError(
                "Anchors should be Tuple[Tuple[int]] because each feature "
                "map could potentially have different sizes and aspect ratios. "
                "There needs to be a match between the number of "
                "feature maps passed and the number of sizes / aspect ratios specified."
            )

        for size, stride, base_anchors in zip(grid_sizes, strides, cell_anchors):
            # for each feature map
            device = base_anchors.device

            # compute anchor centers regarding to the image.
            # shifts_centers is [x_center, y_center] or [x_center, y_center, z_center]
            shifts_centers = [
                torch.arange(0, size[axis], dtype=torch.int32, device=device) * stride[axis]
                for axis in range(self.spatial_dims)
            ]

            # to support torchscript, cannot directly use torch.meshgrid(shifts_centers).
            shifts_centers = list(torch.meshgrid(shifts_centers[: self.spatial_dims], indexing="ij"))

            for axis in range(self.spatial_dims):
                # each element of shifts_centers is sized (HW,) or (HWD,)
                shifts_centers[axis] = shifts_centers[axis].reshape(-1)

            # Expand to [x_center, y_center, x_center, y_center],
            # or [x_center, y_center, z_center, x_center, y_center, z_center]
            if self.indexing == "xy":
                # Cartesian ('xy') indexing swaps axis 0 and 1.
                shifts_centers[1], shifts_centers[0] = shifts_centers[0], shifts_centers[1]
            shifts = torch.stack(shifts_centers * 2, dim=1)  # sized (HW,4) or (HWD,6)

            # For every (base anchor, output anchor) pair,
            # offset each zero-centered base anchor by the center of the output anchor.
            anchors.append(
                (shifts.view(-1, 1, self.spatial_dims * 2) + base_anchors.view(1, -1, self.spatial_dims * 2)).reshape(
                    -1, self.spatial_dims * 2
                )  # each element sized (AHWD,4) or (AHWD,6)
            )

        return anchors

    def forward(self, images: Tensor, feature_maps: list[Tensor]) -> list[Tensor]:
        """
        Generate anchor boxes for each image.

        Args:
            images: sized (B, C, W, H) or (B, C, W, H, D)
            feature_maps: for FPN level i, feature_maps[i] is sized (B, C_i, W_i, H_i) or (B, C_i, W_i, H_i, D_i).
                This input argument does not have to be the actual feature maps.
                Any list variable with the same (C_i, W_i, H_i) or (C_i, W_i, H_i, D_i) as feature maps works.

        Return:
            A list with length of B. Each element represents the anchors for this image.
            The B elements are identical.

        Example:
            .. code-block:: python

                images = torch.zeros((3,1,128,128,128))
                feature_maps = [torch.zeros((3,6,64,64,32)), torch.zeros((3,6,32,32,16))]
                anchor_generator(images, feature_maps)
        """
        grid_sizes = [list(feature_map.shape[-self.spatial_dims :]) for feature_map in feature_maps]
        image_size = images.shape[-self.spatial_dims :]
        batchsize = images.shape[0]
        dtype, device = feature_maps[0].dtype, feature_maps[0].device
        strides = [
            [
                torch.tensor(image_size[axis] // g[axis], dtype=torch.int64, device=device)
                for axis in range(self.spatial_dims)
            ]
            for g in grid_sizes
        ]

        self.set_cell_anchors(dtype, device)
        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes, strides)

        anchors_per_image = torch.cat(list(anchors_over_all_feature_maps))
        return [anchors_per_image] * batchsize


class AnchorGeneratorWithAnchorShape(AnchorGenerator):
    """
    Module that generates anchors for a set of feature maps and
    image sizes, inherited from :py:class:`~monai.apps.detection.networks.utils.anchor_utils.AnchorGenerator`

    The module support computing anchors at multiple base anchor shapes
    per feature map.

    ``feature_map_scales`` should have the same number of elements with the number of feature maps.

    base_anchor_shapes can have an arbitrary number of elements.
    For 2D images, each element represents anchor width and height [w,h].
    For 2D images, each element represents anchor width, height, and depth [w,h,d].

    AnchorGenerator will output a set of ``len(base_anchor_shapes)`` anchors
    per spatial location for feature map ``i``.

    Args:
        feature_map_scales: scale of anchors for each feature map, i.e., each output level of
            the feature pyramid network (FPN). ``len(feature_map_scales)`` is the number of feature maps.
            ``scale[i]*base_anchor_shapes`` represents the anchor shapes for feature map ``i``.
        base_anchor_shapes: a sequence which represents several anchor shapes for one feature map.
            For N-D images, it is a Sequence of N value Sequence.
        indexing: choose from {'xy', 'ij'}, optional
            Cartesian ('xy') or matrix ('ij', default) indexing of output.
            Cartesian ('xy') indexing swaps axis 0 and 1, which is the setting inside torchvision.
            matrix ('ij', default) indexing keeps the original axis not changed.
            See also indexing in https://pytorch.org/docs/stable/generated/torch.meshgrid.html

    Example:
        .. code-block:: python

            # 2D example inputs for a 2-level feature maps
            feature_map_scales = (1, 2)
            base_anchor_shapes = ((10, 10), (6, 12), (12, 6))
            anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes)

            # 3D example inputs for a 2-level feature maps
            feature_map_scales = (1, 2)
            base_anchor_shapes = ((10, 10, 10), (12, 12, 8), (10, 10, 6), (16, 16, 10))
            anchor_generator = AnchorGeneratorWithAnchorShape(feature_map_scales, base_anchor_shapes)
    """

    __annotations__ = {"cell_anchors": list[torch.Tensor]}

    def __init__(
        self,
        feature_map_scales: Sequence[int] | Sequence[float] = (1, 2, 4, 8),
        base_anchor_shapes: Sequence[Sequence[int]] | Sequence[Sequence[float]] = (
            (32, 32, 32),
            (48, 20, 20),
            (20, 48, 20),
            (20, 20, 48),
        ),
        indexing: str = "ij",
    ) -> None:
        nn.Module.__init__(self)

        spatial_dims = len(base_anchor_shapes[0])
        spatial_dims = look_up_option(spatial_dims, [2, 3])
        self.spatial_dims = spatial_dims

        self.indexing = look_up_option(indexing, ["ij", "xy"])

        base_anchor_shapes_t = torch.Tensor(base_anchor_shapes)
        self.cell_anchors = [self.generate_anchors_using_shape(s * base_anchor_shapes_t) for s in feature_map_scales]

    @staticmethod
    def generate_anchors_using_shape(
        anchor_shapes: torch.Tensor, dtype: torch.dtype = torch.float32, device: torch.device | None = None
    ) -> torch.Tensor:
        """
        Compute cell anchor shapes at multiple sizes and aspect ratios for the current feature map.

        Args:
            anchor_shapes: [w, h] or [w, h, d], sized (N, spatial_dims),
                represents N anchor shapes for the current feature map.
            dtype: target data type of the output Tensor.
            device: target device to put the output Tensor data.

        Returns:
            For 2D images, returns [-w/2, -h/2, w/2, h/2];
            For 3D images, returns [-w/2, -h/2, -d/2, w/2, h/2, d/2]
        """
        half_anchor_shapes = anchor_shapes / 2.0
        base_anchors = torch.cat([-half_anchor_shapes, half_anchor_shapes], dim=1)
        return base_anchors.round().to(dtype=dtype, device=device)
