from typing import Any, Dict, List, Optional, Union, cast

from torch import Tensor, stack

from kornia.augmentation import random_generator as rg
from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D

_planckian_coeffs = {
    'blackbody': Tensor(
        [
            [0.6743, 0.4029, 0.0013],
            [0.6281, 0.4241, 0.1665],
            [0.5919, 0.4372, 0.2513],
            [0.5623, 0.4457, 0.3154],
            [0.5376, 0.4515, 0.3672],
            [0.5163, 0.4555, 0.4103],
            [0.4979, 0.4584, 0.4468],
            [0.4816, 0.4604, 0.4782],
            [0.4672, 0.4619, 0.5053],
            [0.4542, 0.4630, 0.5289],
            [0.4426, 0.4638, 0.5497],
            [0.4320, 0.4644, 0.5681],
            [0.4223, 0.4648, 0.5844],
            [0.4135, 0.4651, 0.5990],
            [0.4054, 0.4653, 0.6121],
            [0.3980, 0.4654, 0.6239],
            [0.3911, 0.4655, 0.6346],
            [0.3847, 0.4656, 0.6444],
            [0.3787, 0.4656, 0.6532],
            [0.3732, 0.4656, 0.6613],
            [0.3680, 0.4655, 0.6688],
            [0.3632, 0.4655, 0.6756],
            [0.3586, 0.4655, 0.6820],
            [0.3544, 0.4654, 0.6878],
            [0.3503, 0.4653, 0.6933],
        ]
    ),
    'CIED': Tensor(
        [
            [0.5829, 0.4421, 0.2288],
            [0.5510, 0.4514, 0.2948],
            [0.5246, 0.4576, 0.3488],
            [0.5021, 0.4618, 0.3941],
            [0.4826, 0.4646, 0.4325],
            [0.4654, 0.4667, 0.4654],
            [0.4502, 0.4681, 0.4938],
            [0.4364, 0.4692, 0.5186],
            [0.4240, 0.4700, 0.5403],
            [0.4127, 0.4705, 0.5594],
            [0.4023, 0.4709, 0.5763],
            [0.3928, 0.4713, 0.5914],
            [0.3839, 0.4715, 0.6049],
            [0.3757, 0.4716, 0.6171],
            [0.3681, 0.4717, 0.6281],
            [0.3609, 0.4718, 0.6380],
            [0.3543, 0.4719, 0.6472],
            [0.3480, 0.4719, 0.6555],
            [0.3421, 0.4719, 0.6631],
            [0.3365, 0.4719, 0.6702],
            [0.3313, 0.4719, 0.6766],
            [0.3263, 0.4719, 0.6826],
            [0.3217, 0.4719, 0.6882],
        ]
    ),
}

_planckian_coeffs_ratio = {
    'blackbody': stack(
        (
            _planckian_coeffs['blackbody'][:, 0] / _planckian_coeffs['blackbody'][:, 1],
            _planckian_coeffs['blackbody'][:, 2] / _planckian_coeffs['blackbody'][:, 1],
        ),
        1,
    ),
    'CIED': stack(
        (
            _planckian_coeffs['CIED'][:, 0] / _planckian_coeffs['CIED'][:, 1],
            _planckian_coeffs['CIED'][:, 2] / _planckian_coeffs['CIED'][:, 1],
        ),
        1,
    ),
}


class RandomPlanckianJitter(IntensityAugmentationBase2D):
    r"""Apply planckian jitter transformation to input tensor.

    .. image:: _static/img/RandomPlanckianJitter.png

    This is physics based color augmentation, that creates realistic
    variations in chromaticity, this can simulate the illumination
    changes in the scene.

    See :cite:`zini2022planckian` for more details.

    Args:
        mode: 'blackbody' or 'CIED'.
        select_from: choose a list of jitters to apply from. `blackbody` range [0-24], `CIED` range [0-22]
        same_on_batch: apply the same transformation across the batch.
        p: probability that the random erasing operation will be performed.
        keepdim: whether to keep the output shape the same as input (True) or broadcast it
            to the batch form (False).

    Shape:
        - Input: :math:`(C, H, W)` or :math:`(B, C, H, W)`
        - Output: :math:`(B, C, H, W)`

    .. note::
        Input tensor must be float and normalized into [0, 1].

    Examples:

        To apply planckian jitter based on mode

        >>> rng = torch.manual_seed(0)
        >>> input = torch.randn(1, 3, 2, 2)
        >>> aug = RandomPlanckianJitter(mode='CIED')
        >>> aug(input)
        tensor([[[[ 1.0000, -0.2389],
                  [-1.7740,  0.4628]],
        <BLANKLINE>
                 [[-1.0845, -1.3986],
                  [ 0.4033,  0.8380]],
        <BLANKLINE>
                 [[-0.9228, -0.5175],
                  [-0.7654,  0.2335]]]])

        To apply planckian jitter on image(s) from list of interested jitters

        >>> rng = torch.manual_seed(0)
        >>> input = torch.randn(2, 3, 2, 2)
        >>> aug = RandomPlanckianJitter(mode='blackbody', select_from=[23, 24, 1, 2])
        >>> aug(input)
        tensor([[[[-1.1258, -1.1524],
                  [-0.2506, -0.4339]],
        <BLANKLINE>
                 [[ 0.8487,  0.6920],
                  [-0.3160, -2.1152]],
        <BLANKLINE>
                 [[ 0.4681, -0.1577],
                  [ 1.4437,  0.2660]]],
        <BLANKLINE>
        <BLANKLINE>
                [[[ 0.2465,  1.0000],
                  [-0.2125, -0.1653]],
        <BLANKLINE>
                 [[ 0.9318,  1.0000],
                  [ 1.0000,  0.0537]],
        <BLANKLINE>
                 [[ 0.2426, -0.1621],
                  [-0.3302, -0.9093]]]])
    """

    def __init__(
        self,
        mode: str = "blackbody",
        select_from: Optional[Union[int, List[int]]] = None,
        same_on_batch: bool = False,
        p: float = 0.5,
        keepdim: bool = False,
        return_transform: Optional[bool] = None,
    ) -> None:
        super().__init__(p=p, return_transform=return_transform, same_on_batch=same_on_batch, keepdim=keepdim)

        if isinstance(select_from, int):
            select_from = [select_from]

        if select_from is not None:
            self.register_buffer('pl', _planckian_coeffs_ratio[mode][select_from])
        else:
            self.register_buffer('pl', _planckian_coeffs_ratio[mode])

        # the range of the sampling parameters
        _param_min: float = 0.0
        _param_max: float = float(cast(Tensor, self.pl).shape[0])

        self._param_generator = cast(rg.PlanckianJitterGenerator, rg.PlanckianJitterGenerator([_param_min, _param_max]))

    def apply_transform(
        self, input: Tensor, params: Dict[str, Tensor], flags: Dict[str, Any], transform: Optional[Tensor] = None
    ) -> Tensor:

        list_idx = params['idx'].tolist()
        coeffs = cast(Tensor, self.pl)[list_idx]

        r_w = coeffs[:, 0][..., None, None]
        b_w = coeffs[:, 1][..., None, None]

        r = input[..., 0, :, :] * r_w
        g = input[..., 1, :, :]
        b = input[..., 2, :, :] * b_w

        output = stack([r, g, b], -3)

        return output.clamp(max=1.0)