# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence

from mmeval.core import BaseMetric
from mmeval.utils import try_import
from .utils import reorder_and_crop

if TYPE_CHECKING:
    import cv2
else:
    cv2 = try_import('cv2')


class SSIM(BaseMetric):
    """Calculate SSIM (structural similarity).

    Ref:
    Image quality assessment: From error visibility to structural similarity

    The results are the same as that of the official released MATLAB code in
    https://ece.uwaterloo.ca/~z70wang/research/ssim/.

    For three-channel images, SSIM is calculated for each channel and then
    averaged.

    Args:
        crop_border (int): Cropped pixels in each edges of an image. These
            pixels are not involved in the PSNR calculation. Defaults to 0.
        input_order (str): Whether the input order is 'HWC' or 'CHW'.
            Defaults to 'HWC'.
        convert_to (str, optional): Whether to convert the images to other
            color models. If None, the images are not altered. When computing
            for 'Y', the images are assumed to be in BGR order. Options are
            'Y' and None. Defaults to None.
        channel_order (str): The channel order of image. Choices are 'rgb' and
            'bgr'. Defaults to 'rgb'.
        **kwargs: Keyword parameters passed to :class:`BaseMetric`.

    Examples:

        >>> from mmeval import SSIM
        >>> import numpy as np
        >>>
        >>> ssim = SSIM(input_order='CHW', convert_to='Y', channel_order='rgb')
        >>> gts = np.random.randint(0, 255, size=(3, 32, 32))
        >>> preds = np.random.randint(0, 255, size=(3, 32, 32))
        >>> ssim(preds, gts)  # doctest: +ELLIPSIS
        {'ssim': ...}

    Calculate SSIM between 2 single channel images:

        >>> img1 = np.ones((32, 32)) * 2
        >>> img2 = np.ones((32, 32))
        >>> SSIM.compute_ssim(img1, img2)
        0.913062377743969
    """

    def __init__(self,
                 crop_border: int = 0,
                 input_order: str = 'CHW',
                 convert_to: Optional[str] = None,
                 channel_order: str = 'rgb',
                 **kwargs) -> None:
        super().__init__(**kwargs)

        assert input_order.upper() in [
            'CHW', 'HWC'
        ], (f'Wrong input_order {input_order}. Supported input_orders are '
            '"HWC" and "CHW"')
        self.input_order = input_order
        self.crop_border = crop_border

        if convert_to is not None:
            assert convert_to.upper() == 'Y', (
                'Wrong color model. Supported values are "Y" and None.')
            assert channel_order.upper() in [
                'BGR', 'RGB'
            ], ('Only support `rgb2y` and `bgr2y`, but the channel_order '
                f'is {channel_order}')
        self.channel_order = channel_order
        self.convert_to = convert_to

    def add(self, predictions: Sequence[np.ndarray], groundtruths: Sequence[np.ndarray], channel_order: Optional[str] = None) -> None:   # type: ignore # yapf: disable # noqa: E501
        """Add the SSIM score of the batch to ``self._results``.

        For three-channel images, SSIM is calculated for each channel and then
        averaged.

        Args:
            predictions (Sequence[np.ndarray]): Predictions of the model.
            groundtruths (Sequence[np.ndarray]): The ground truth images.
            channel_order (Optional[str]): The channel order of the input
                samples. If not passed, will set as :attr:`self.channel_order`.
                Defaults to None.
        """
        channel_order = self.channel_order \
            if channel_order is None else channel_order

        for pred, gt in zip(predictions, groundtruths):
            assert pred.shape == gt.shape, (
                f'Image shapes are different: {pred.shape}, {gt.shape}.')
            pred = reorder_and_crop(
                pred,
                crop_border=self.crop_border,
                input_order=self.input_order,
                convert_to=self.convert_to,
                channel_order=channel_order)
            gt = reorder_and_crop(
                gt,
                crop_border=self.crop_border,
                input_order=self.input_order,
                convert_to=self.convert_to,
                channel_order=channel_order)

            _ssim_score = []
            for i in range(pred.shape[2]):
                _ssim_score.append(self.compute_ssim(pred[..., i], gt[..., i]))
            self._results.append(np.array(_ssim_score).mean())

    def compute_metric(self, results: List[np.float64]) -> Dict[str, float]:
        """Compute the SSIM metric.

        This method would be invoked in ``BaseMetric.compute`` after
        distributed synchronization.

        Args:
            results (List[np.float64]): A list that consisting the SSIM scores.
                This list has already been synced across all ranks.

        Returns:
            Dict[str, float]: The computed SSIM metric.
        """

        return {'ssim': float(np.array(results).mean())}

    @staticmethod
    def compute_ssim(img1: np.ndarray, img2: np.ndarray) -> np.float64:
        """Calculate SSIM (structural similarity) between two single channel
        image.

        Ref:
        Image quality assessment: From error visibility to structural
        similarity

        The results are the same as that of the official released MATLAB code
        in https://ece.uwaterloo.ca/~z70wang/research/ssim/.

        Args:
            img1 (np.ndarray): Single channels Images with range [0, 255].
            img2 (np.ndarray): Single channels Images with range [0, 255].

        Returns:
            np.float64: SSIM result.
        """

        C1 = (0.01 * 255)**2
        C2 = (0.03 * 255)**2

        kernel = cv2.getGaussianKernel(11, 1.5)
        window = np.outer(kernel, kernel.transpose())

        mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]
        mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
        mu1_sq = mu1**2
        mu2_sq = mu2**2
        mu1_mu2 = mu1 * mu2
        sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
        sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
        sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

        ssim_map = ((2 * mu1_mu2 + C1) *
                    (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
                                           (sigma1_sq + sigma2_sq + C2))

        return ssim_map.mean()
