#!/usr/bin/env python3
""" Common multi-backend Keras utilities """
from __future__ import annotations
import typing as T

import numpy as np

from keras import ops, Variable

from lib.utils import get_module_objects

if T.TYPE_CHECKING:
    from keras import KerasTensor

# TODO these can probably be switched to pure pytorch


def frobenius_norm(matrix: KerasTensor,
                   axis: int = -1,
                   keep_dims: bool = True,
                   epsilon: float = 1e-15) -> KerasTensor:
    """ Frobenius normalization for Keras Tensor

    Parameters
    ----------
    matrix: :class:`keras.KerasTensor`
        The matrix to normalize
    axis: int, optional
        The axis to normalize. Default: `-1`
    keep_dims: bool, Optional
        Whether to retain the original matrix shape or not. Default:``True``
    epsilon: flot, optional
        Epsilon to apply to the normalization to preven NaN errors on zero values

    Returns
    -------
    :class:`keras.KerasTensor`
        The normalized output
    """
    return ops.sqrt(ops.sum(ops.power(matrix, 2), axis=axis, keepdims=keep_dims) + epsilon)


def replicate_pad(image: KerasTensor, padding: int) -> KerasTensor:
    """ Apply replication padding to an input batch of images. Expects 4D tensor in BHWC format.

    Notes
    -----
    At the time of writing Keras does not have a native replication padding method.
    The implementation here is probably not the most efficient, but it is a pure keras method
    which should work ok.

    Parameters
    ----------
    image: :class:`keras.KerasTensor`
        Image tensor to pad
    pad: int
        The amount of padding to apply to each side of the input image

    Returns
    -------
    :class:`keras.KerasTensor`
        The input image with replication padding applied
    """
    top_pad = ops.tile(image[:, :1, ...], (1, padding, 1, 1))
    bottom_pad = ops.tile(image[:, -1:, ...], (1, padding, 1, 1))
    pad_top_bottom = ops.concatenate([top_pad, image, bottom_pad], axis=1)
    left_pad = ops.tile(pad_top_bottom[..., :1, :], (1, 1, padding, 1))
    right_pad = ops.tile(pad_top_bottom[..., -1:, :],  (1, 1, padding, 1))
    padded = ops.concatenate([left_pad, pad_top_bottom, right_pad], axis=2)
    return padded


class ColorSpaceConvert():
    """ Transforms inputs between different color spaces on the GPU

    Notes
    -----
    The following color space transformations are implemented:
        - rgb to lab
        - rgb to xyz
        - srgb to _rgb
        - srgb to ycxcz
        - xyz to ycxcz
        - xyz to lab
        - xyz to rgb
        - ycxcz to rgb
        - ycxcz to xyz

    Parameters
    ----------
    from_space: str
        One of `"srgb"`, `"rgb"`, `"xyz"`
    to_space: str
        One of `"lab"`, `"rgb"`, `"ycxcz"`, `"xyz"`

    Raises
    ------
    ValueError
        If the requested color space conversion is not defined
    """
    def __init__(self, from_space: str, to_space: str) -> None:
        functions = {"rgb_lab": self._rgb_to_lab,
                     "rgb_xyz": self._rgb_to_xyz,
                     "srgb_rgb": self._srgb_to_rgb,
                     "srgb_ycxcz": self._srgb_to_ycxcz,
                     "xyz_ycxcz": self._xyz_to_ycxcz,
                     "xyz_lab": self._xyz_to_lab,
                     "xyz_rgb": self._xyz_to_rgb,
                     "ycxcz_rgb": self._ycxcz_to_rgb,
                     "ycxcz_xyz": self._ycxcz_to_xyz}
        func_name = f"{from_space.lower()}_{to_space.lower()}"
        if func_name not in functions:
            raise ValueError(f"The color transform {from_space} to {to_space} is not defined.")

        self._func = functions[func_name]
        self._ref_illuminant = Variable(np.array([[[0.950428545, 1.000000000, 1.088900371]]]),
                                        dtype="float32",
                                        trainable=False)
        self._inv_ref_illuminant = 1. / self._ref_illuminant

        self._rgb_xyz_map = self._get_rgb_xyz_map()
        self._xyz_multipliers = Variable([116, 500, 200], dtype="float32", trainable=False)

    @classmethod
    def _get_rgb_xyz_map(cls) -> tuple[KerasTensor, KerasTensor]:
        """ Obtain the mapping and inverse mapping for rgb to xyz color space conversion.

        Returns
        -------
        tuple
            The mapping and inverse Tensors for rgb to xyz color space conversion
        """
        mapping = np.array([[10135552 / 24577794,  8788810 / 24577794, 4435075 / 24577794],
                            [2613072 / 12288897, 8788810 / 12288897, 887015 / 12288897],
                            [1425312 / 73733382, 8788810 / 73733382, 70074185 / 73733382]])
        inverse = np.linalg.inv(mapping)
        return (Variable(mapping, dtype="float32", trainable=False),
                Variable(inverse, dtype="float32", trainable=False))

    def __call__(self, image: KerasTensor) -> KerasTensor:
        """ Call the colorspace conversion function.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in the colorspace defined by :attr:`from_space`

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in the colorspace defined by :attr:`to_space`
        """
        return self._func(image)

    def _rgb_to_lab(self, image: KerasTensor) -> KerasTensor:
        """ RGB to LAB conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in RGB format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in LAB format
        """
        converted = self._rgb_to_xyz(image)
        return self._xyz_to_lab(converted)

    def _rgb_xyz_rgb(self, image: KerasTensor, mapping: KerasTensor) -> KerasTensor:
        """ RGB to XYZ or XYZ to RGB conversion.

        Notes
        -----
        The conversion in both directions is the same, but the mappping matrix for XYZ to RGB is
        the inverse of RGB to XYZ.

        References
        ----------
        https://www.image-engineering.de/library/technotes/958-how-to-convert-between-srgb-and-ciexyz

        Parameters
        ----------
        mapping: :class:`keras.KerasTensor`
            The mapping matrix to perform either the XYZ to RGB or RGB to XYZ color space
            conversion

        image: :class:`keras.KerasTensor`
            The image tensor in RGB format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in XYZ format
        """
        dim = image.shape
        image = ops.transpose(image, (0, 3, 1, 2))
        image = ops.reshape(image, (dim[0], dim[3], dim[1] * dim[2]))
        converted = ops.transpose(ops.dot(mapping, image), (0, 2, 1))
        return ops.reshape(converted, dim)

    def _rgb_to_xyz(self, image: KerasTensor) -> KerasTensor:
        """ RGB to XYZ conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in RGB format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in XYZ format
        """
        return self._rgb_xyz_rgb(image, self._rgb_xyz_map[0])

    @classmethod
    def _srgb_to_rgb(cls, image: KerasTensor) -> KerasTensor:
        """ SRGB to RGB conversion.

        Notes
        -----
        RGB Image is clipped to a small epsilon to stabalize training

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in SRGB format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in RGB format
        """
        limit = np.float32(0.04045)
        return ops.where(image > limit,
                         ops.power((ops.clip(image, limit, np.inf) + 0.055) / 1.055, 2.4),
                         image / 12.92)

    def _srgb_to_ycxcz(self, image: KerasTensor) -> KerasTensor:
        """ SRGB to YcXcZ conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in SRGB format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in YcXcZ format
        """
        converted = self._srgb_to_rgb(image)
        converted = self._rgb_to_xyz(converted)
        return self._xyz_to_ycxcz(converted)

    def _xyz_to_lab(self, image: KerasTensor) -> KerasTensor:
        """ XYZ to LAB conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in XYZ format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in LAB format
        """
        image = image * self._inv_ref_illuminant
        delta = 6 / 29
        delta_cube = delta ** 3
        factor = 1 / (3 * (delta ** 2))

        clamped_term = ops.power(ops.clip(image, delta_cube, np.inf), 1.0 / 3.0)
        div = factor * image + (4 / 29)

        image = ops.where(image > delta_cube, clamped_term, div)

        return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16.,
                                self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])],
                               axis=-1)

    def _xyz_to_rgb(self, image: KerasTensor) -> KerasTensor:
        """ XYZ to YcXcZ conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in XYZ format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in RGB format
        """
        return self._rgb_xyz_rgb(image, self._rgb_xyz_map[1])

    def _xyz_to_ycxcz(self, image: KerasTensor) -> KerasTensor:
        """ XYZ to YcXcZ conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in XYZ format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in YcXcZ format
        """
        image = image * self._inv_ref_illuminant
        return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16.,
                                self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])],
                               axis=-1)

    def _ycxcz_to_rgb(self, image: KerasTensor) -> KerasTensor:
        """ YcXcZ to RGB conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in YcXcZ format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in RGB format
        """
        converted = self._ycxcz_to_xyz(image)
        return self._xyz_to_rgb(converted)

    def _ycxcz_to_xyz(self, image: KerasTensor) -> KerasTensor:
        """ YcXcZ to XYZ conversion.

        Parameters
        ----------
        image: :class:`keras.KerasTensor`
            The image tensor in YcXcZ format

        Returns
        -------
        :class:`keras.KerasTensor`
            The image tensor in XYZ format
        """
        ch_y = (image[..., 0:1] + 16.) / self._xyz_multipliers[0]
        return ops.concatenate([ch_y + (image[..., 1:2] / self._xyz_multipliers[1]),
                                ch_y,
                                ch_y - (image[..., 2:3] / self._xyz_multipliers[2])],
                               axis=-1) * self._ref_illuminant


__all__ = get_module_objects(__name__)
