"""Functional implementations for domain adaptation image transformations.

This module provides low-level functions and classes for performing domain adaptation
between images. It includes implementations for histogram matching, Fourier domain adaptation,
and pixel distribution matching with various normalization techniques.
"""

import abc
from copy import deepcopy
from typing import Literal

import cv2
import numpy as np
from albucore import add_weighted, clip, clipped, from_float, get_num_channels, preserve_channel_dim, to_float, uint8_io
from typing_extensions import Protocol

import albumentations.augmentations.geometric.functional as fgeometric
from albumentations.augmentations.utils import PCA
from albumentations.core.type_definitions import MONO_CHANNEL_DIMENSIONS, ImageType

__all__ = [
    "adapt_pixel_distribution",
    "apply_histogram",
    "fourier_domain_adaptation",
]


class BaseScaler:
    def __init__(self) -> None:
        self.data_min: np.ndarray | None = None
        self.data_max: np.ndarray | None = None
        self.mean: np.ndarray | None = None
        self.var: np.ndarray | None = None
        self.scale: np.ndarray | None = None

    def fit(self, x: np.ndarray) -> None:
        """Fit the scaler to the data.

        Args:
            x (np.ndarray): The data to fit the scaler to. Expected shape is (n_samples, n_features).

        Raises:
            NotImplementedError: This is an abstract method that must be implemented by subclasses.

        """
        raise NotImplementedError

    def transform(self, x: np.ndarray) -> np.ndarray:
        """Transform the data using the fitted scaler.

        Args:
            x (np.ndarray): The data to transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The transformed data with the same shape as the input.

        Raises:
            NotImplementedError: This is an abstract method that must be implemented by subclasses.

        """
        raise NotImplementedError

    def fit_transform(self, x: np.ndarray) -> np.ndarray:
        """Fit the scaler to the data and then transform it.

        This is a convenience method that combines fit and transform operations.

        Args:
            x (np.ndarray): The data to fit and transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The transformed data with the same shape as the input.

        """
        self.fit(x)
        return self.transform(x)

    def inverse_transform(self, x: np.ndarray) -> np.ndarray:
        """Apply the inverse transformation to the data.

        This method reverses the transformation applied by transform().

        Args:
            x (np.ndarray): The transformed data to inverse transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The inverse transformed data with the same shape as the input.

        Raises:
            NotImplementedError: This is an abstract method that must be implemented by subclasses.

        """
        raise NotImplementedError


class MinMaxScaler(BaseScaler):
    def __init__(self, feature_range: tuple[float, float] = (0.0, 1.0)) -> None:
        super().__init__()
        self.min: float = feature_range[0]
        self.max: float = feature_range[1]
        self.data_range: np.ndarray | None = None

    def fit(self, x: np.ndarray) -> None:
        """Fit the MinMaxScaler to the data.

        Computes the minimum and maximum values for each feature to be used for scaling.

        Args:
            x (np.ndarray): The data to fit the scaler to. Expected shape is (n_samples, n_features).

        Note:
            When data_min equals data_max for a feature, the range is set to 1 to avoid division by zero.

        """
        self.data_min = np.min(x, axis=0)
        self.data_max = np.max(x, axis=0)
        self.data_range = self.data_max - self.data_min
        # Handle case where data_min equals data_max
        self.data_range[self.data_range == 0] = 1

    def transform(self, x: np.ndarray) -> np.ndarray:
        """Transform features by scaling to the specified range.

        The transformation is: X_scaled = (X - data_min) / data_range * (max - min) + min

        Args:
            x (np.ndarray): The data to transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The scaled data in the range [self.min, self.max].

        Raises:
            ValueError: If the scaler has not been fitted yet.

        """
        if self.data_min is None or self.data_max is None or self.data_range is None:
            raise ValueError(
                "This MinMaxScaler instance is not fitted yet. "
                "Call 'fit' with appropriate arguments before using this estimator.",
            )

        x_std = np.subtract(x, self.data_min).astype(float)
        np.divide(x_std, self.data_range, out=x_std)
        np.multiply(x_std, (self.max - self.min), out=x_std)
        np.add(x_std, self.min, out=x_std)

        return x_std

    def inverse_transform(self, x: np.ndarray) -> np.ndarray:
        """Reverse the scaling operation.

        Args:
            x (np.ndarray): The scaled data to inverse transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The data in the original scale.

        Raises:
            ValueError: If the scaler has not been fitted yet.

        """
        if self.data_min is None or self.data_max is None or self.data_range is None:
            raise ValueError(
                "This MinMaxScaler instance is not fitted yet. "
                "Call 'fit' with appropriate arguments before using this estimator.",
            )
        x_std = ((x - self.min) / (self.max - self.min)).astype(float)
        return x_std * self.data_range + self.data_min


class StandardScaler(BaseScaler):
    def __init__(self) -> None:
        super().__init__()

    def fit(self, x: np.ndarray) -> None:
        """Fit the StandardScaler to the data.

        Computes the mean and standard deviation for each feature to be used for standardization.

        Args:
            x (np.ndarray): The data to fit the scaler to. Expected shape is (n_samples, n_features).

        Note:
            When variance is zero for a feature, the scale is set to 1 to avoid division by zero.

        """
        self.mean = np.mean(x, axis=0)
        self.var = np.var(x, axis=0)
        self.scale = np.sqrt(self.var)
        # Handle case where variance is zero
        self.scale[self.scale == 0] = 1

    def transform(self, x: np.ndarray) -> np.ndarray:
        """Standardize features by removing the mean and scaling to unit variance.

        The transformation is: z = (x - mean) / std

        Args:
            x (np.ndarray): The data to transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The standardized data with zero mean and unit variance.

        Raises:
            ValueError: If the scaler has not been fitted yet.

        """
        if self.mean is None or self.scale is None:
            raise ValueError(
                "This StandardScaler instance is not fitted yet. "
                "Call 'fit' with appropriate arguments before using this estimator.",
            )
        return (x - self.mean) / self.scale

    def inverse_transform(self, x: np.ndarray) -> np.ndarray:
        """Reverse the standardization operation.

        The inverse transformation is: x = z * std + mean

        Args:
            x (np.ndarray): The standardized data to inverse transform. Expected shape is (n_samples, n_features).

        Returns:
            np.ndarray: The data in the original scale.

        Raises:
            ValueError: If the scaler has not been fitted yet.

        """
        if self.mean is None or self.scale is None:
            raise ValueError(
                "This StandardScaler instance is not fitted yet. "
                "Call 'fit' with appropriate arguments before using this estimator.",
            )
        return (x * self.scale) + self.mean


class TransformerInterface(Protocol):
    @abc.abstractmethod
    def inverse_transform(self, x: np.ndarray) -> np.ndarray:
        """Apply the inverse transformation.

        Args:
            x (np.ndarray): The transformed data to inverse transform.

        Returns:
            np.ndarray: The inverse transformed data.

        """
        ...

    @abc.abstractmethod
    def fit(self, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
        """Fit the transformer to the data.

        Args:
            x (np.ndarray): The data to fit to.
            y (np.ndarray | None): Optional target data (not used in most implementations).

        Returns:
            np.ndarray: The fitted transformer instance.

        """
        ...

    @abc.abstractmethod
    def transform(self, x: np.ndarray, y: np.ndarray | None = None) -> np.ndarray:
        """Transform the data.

        Args:
            x (np.ndarray): The data to transform.
            y (np.ndarray | None): Optional target data (not used in most implementations).

        Returns:
            np.ndarray: The transformed data.

        """
        ...


class DomainAdapter:
    def __init__(
        self,
        transformer: TransformerInterface,
        ref_img: ImageType,
        color_conversions: tuple[None, None] = (None, None),
    ):
        self.color_in, self.color_out = color_conversions
        self.source_transformer = deepcopy(transformer)
        self.target_transformer = transformer
        self.num_channels = get_num_channels(ref_img)
        self.target_transformer.fit(self.flatten(ref_img))

    def to_colorspace(self, img: ImageType) -> ImageType:
        """Convert the image to the target color space.

        Args:
            img (np.ndarray): The input image to convert.

        Returns:
            np.ndarray: The image in the target color space, or the original image if no conversion is specified.

        """
        return img if self.color_in is None else cv2.cvtColor(img, self.color_in)

    def from_colorspace(self, img: ImageType) -> ImageType:
        """Convert the image back from the target color space.

        Args:
            img (np.ndarray): The image to convert back.

        Returns:
            np.ndarray: The image converted back to the original color space, or the original image
                if no conversion is specified.

        """
        if self.color_out is None:
            return img
        return cv2.cvtColor(clip(img, np.uint8, inplace=True), self.color_out)

    def flatten(self, img: ImageType) -> np.ndarray:
        """Flatten the image into a 2D array of pixels.

        Converts the image to the target color space, normalizes to float values,
        and reshapes to (n_pixels, n_channels).

        Args:
            img (np.ndarray): The input image to flatten.

        Returns:
            np.ndarray: The flattened image with shape (n_pixels, n_channels).

        """
        img = self.to_colorspace(img)
        img = to_float(img)
        return img.reshape(-1, self.num_channels)

    def reconstruct(self, pixels: np.ndarray, height: int, width: int) -> np.ndarray:
        """Reconstruct an image from flattened pixels.

        Reshapes the pixels back to image format and converts back to the original color space.

        Args:
            pixels (np.ndarray): The flattened pixels with shape (n_pixels, n_channels).
            height (int): The height of the output image.
            width (int): The width of the output image.

        Returns:
            np.ndarray: The reconstructed image with shape (height, width) for grayscale
                or (height, width, n_channels) for color images.

        """
        pixels = clip(pixels, np.uint8, inplace=True)
        if self.num_channels == 1:
            return self.from_colorspace(pixels.reshape(height, width))
        return self.from_colorspace(pixels.reshape(height, width, self.num_channels))

    @staticmethod
    def _pca_sign(x: np.ndarray) -> np.ndarray:
        return np.sign(np.trace(x.components_))

    def __call__(self, image: ImageType) -> ImageType:
        height, width = image.shape[:2]
        pixels = self.flatten(image)
        self.source_transformer.fit(pixels)

        if (
            hasattr(self.target_transformer, "components_")
            and hasattr(self.source_transformer, "components_")
            and self._pca_sign(self.target_transformer) != self._pca_sign(self.source_transformer)
        ):
            self.target_transformer.components_ *= -1

        representation = self.source_transformer.transform(pixels)
        result = self.target_transformer.inverse_transform(representation)
        return self.reconstruct(result, height, width)


@clipped
@preserve_channel_dim
def adapt_pixel_distribution(
    img: ImageType,
    ref: np.ndarray,
    transform_type: Literal["pca", "standard", "minmax"],
    weight: float,
) -> ImageType:
    """Adapt the pixel distribution of an image to match a reference image.

    This function adapts the pixel distribution of an image to match a reference image
    using a specified transformation type and weight.

    Args:
        img (np.ndarray): The input image to be adapted.
        ref (np.ndarray): The reference image.
        transform_type (Literal["pca", "standard", "minmax"]): The type of transformation to use.
        weight (float): The weight of the transformation.

    Returns:
        np.ndarray: The adapted image.

    Raises:
        ValueError: If the input image and reference image have different dtypes or numbers of channels.

    """
    if img.dtype != ref.dtype:
        raise ValueError("Input image and reference image must have the same dtype.")
    img_num_channels = get_num_channels(img)
    ref_num_channels = get_num_channels(ref)

    if img_num_channels != ref_num_channels:
        raise ValueError("Input image and reference image must have the same number of channels.")

    if img_num_channels == 1:
        img = np.squeeze(img)
        ref = np.squeeze(ref)

    if img.shape != ref.shape:
        ref = cv2.resize(ref, dsize=img.shape[:2], interpolation=cv2.INTER_AREA)

    original_dtype = img.dtype

    if original_dtype == np.float32:
        img = from_float(img, np.uint8)
        ref = from_float(ref, np.uint8)

    transformer = {"pca": PCA, "standard": StandardScaler, "minmax": MinMaxScaler}[transform_type]()
    adapter = DomainAdapter(transformer=transformer, ref_img=ref)
    transformed = adapter(img).astype(np.float32)

    result = img.astype(np.float32) * (1 - weight) + transformed * weight

    return result if original_dtype == np.uint8 else to_float(result)


def low_freq_mutate(amp_src: np.ndarray, amp_trg: np.ndarray, beta: float) -> np.ndarray:
    image_shape = amp_src.shape[:2]

    border = int(np.floor(min(image_shape) * beta))

    center_x, center_y = fgeometric.center(image_shape)

    height, width = image_shape

    h1, h2 = max(0, int(center_y - border)), min(int(center_y + border), height)
    w1, w2 = max(0, int(center_x - border)), min(int(center_x + border), width)
    amp_src[h1:h2, w1:w2] = amp_trg[h1:h2, w1:w2]
    return amp_src


@clipped
@preserve_channel_dim
def fourier_domain_adaptation(img: ImageType, target_img: ImageType, beta: float) -> ImageType:
    """Apply Fourier Domain Adaptation to the input image using a target image.

    This function performs domain adaptation in the frequency domain by modifying the amplitude
    spectrum of the source image based on the target image's amplitude spectrum. It preserves
    the phase information of the source image, which helps maintain its content while adapting
    its style to match the target image.

    Args:
        img (np.ndarray): The source image to be adapted. Can be grayscale or RGB.
        target_img (np.ndarray): The target image used as a reference for adaptation.
            Should have the same dimensions as the source image.
        beta (float): The adaptation strength, typically in the range [0, 1].
            Higher values result in stronger adaptation towards the target image's style.

    Returns:
        np.ndarray: The adapted image with the same shape and type as the input image.

    Raises:
        ValueError: If the source and target images have different shapes.

    Note:
        - Both input images are converted to float32 for processing.
        - The function handles both grayscale (2D) and color (3D) images.
        - For grayscale images, an extra dimension is added to facilitate uniform processing.
        - The adaptation is performed channel-wise for color images.
        - The output is clipped to the valid range and preserves the original number of channels.

    The adaptation process involves the following steps for each channel:
    1. Compute the 2D Fourier Transform of both source and target images.
    2. Shift the zero frequency component to the center of the spectrum.
    3. Extract amplitude and phase information from the source image's spectrum.
    4. Mutate the source amplitude using the target amplitude and the beta parameter.
    5. Combine the mutated amplitude with the original phase.
    6. Perform the inverse Fourier Transform to obtain the adapted channel.

    The `low_freq_mutate` function (not shown here) is responsible for the actual
    amplitude mutation, focusing on low-frequency components which carry style information.

    Examples:
        >>> import numpy as np
        >>> import albumentations as A
        >>> source_img = np.random.rand(100, 100, 3).astype(np.float32)
        >>> target_img = np.random.rand(100, 100, 3).astype(np.float32)
        >>> adapted_img = A.fourier_domain_adaptation(source_img, target_img, beta=0.5)
        >>> assert adapted_img.shape == source_img.shape

    References:
        FDA: Fourier Domain Adaptation for Semantic Segmentation: Yang and Soatto, 2020, CVPR
            https://openaccess.thecvf.com/content_CVPR_2020/papers/Yang_FDA_Fourier_Domain_Adaptation_for_Semantic_Segmentation_CVPR_2020_paper.pdf

    """
    src_img = img.astype(np.float32)
    trg_img = target_img.astype(np.float32)

    if src_img.ndim == MONO_CHANNEL_DIMENSIONS:
        src_img = np.expand_dims(src_img, axis=-1)
    if trg_img.ndim == MONO_CHANNEL_DIMENSIONS:
        trg_img = np.expand_dims(trg_img, axis=-1)

    num_channels = src_img.shape[-1]

    # Prepare container for the output image
    src_in_trg = np.zeros_like(src_img)

    for channel_id in range(num_channels):
        # Perform FFT on each channel
        fft_src = np.fft.fft2(src_img[:, :, channel_id])
        fft_trg = np.fft.fft2(trg_img[:, :, channel_id])

        # Shift the zero frequency component to the center
        fft_src_shifted = np.fft.fftshift(fft_src)
        fft_trg_shifted = np.fft.fftshift(fft_trg)

        # Extract amplitude and phase
        amp_src, pha_src = np.abs(fft_src_shifted), np.angle(fft_src_shifted)
        amp_trg = np.abs(fft_trg_shifted)

        # Mutate the amplitude part of the source with the target
        mutated_amp = low_freq_mutate(amp_src.copy(), amp_trg, beta)

        # Combine the mutated amplitude with the original phase
        fft_src_mutated = np.fft.ifftshift(mutated_amp * np.exp(1j * pha_src))

        # Perform inverse FFT
        src_in_trg_channel = np.fft.ifft2(fft_src_mutated)

        # Store the result in the corresponding channel of the output image
        src_in_trg[:, :, channel_id] = np.real(src_in_trg_channel)

    return src_in_trg


@clipped
@preserve_channel_dim
def apply_histogram(img: ImageType, reference_image: ImageType, blend_ratio: float) -> ImageType:
    """Apply histogram matching to an input image using a reference image and blend the result.

    This function performs histogram matching between the input image and a reference image,
    then blends the result with the original input image based on the specified blend ratio.

    Args:
        img (np.ndarray): The input image to be transformed. Can be either grayscale or RGB.
            Supported dtypes: uint8, float32 (values should be in [0, 1] range).
        reference_image (np.ndarray): The reference image used for histogram matching.
            Should have the same number of channels as the input image.
            Supported dtypes: uint8, float32 (values should be in [0, 1] range).
        blend_ratio (float): The ratio for blending the matched image with the original image.
            Should be in the range [0, 1], where 0 means no change and 1 means full histogram matching.

    Returns:
        np.ndarray: The transformed image after histogram matching and blending.
            The output will have the same shape and dtype as the input image.

    Supported image types:
        - Grayscale images: 2D arrays
        - RGB images: 3D arrays with 3 channels
        - Multispectral images: 3D arrays with more than 3 channels

    Note:
        - If the input and reference images have different sizes, the reference image
          will be resized to match the input image's dimensions.
        - The function uses a custom implementation of histogram matching based on OpenCV and NumPy.
        - The @clipped and @preserve_channel_dim decorators ensure the output is within
          the valid range and maintains the original number of dimensions.

    """
    # Resize reference image only if necessary
    if img.shape[:2] != reference_image.shape[:2]:
        reference_image = cv2.resize(reference_image, dsize=(img.shape[1], img.shape[0]))

    reference_image = reference_image.reshape(img.shape)

    # Match histograms between the images
    matched = match_histograms(img, reference_image)

    # Blend the original image and the matched image
    return add_weighted(matched, blend_ratio, img, 1 - blend_ratio)


@uint8_io
def match_histograms(image: ImageType, reference: ImageType) -> ImageType:
    """Adjust an image so that its cumulative histogram matches that of another.

    The adjustment is applied separately for each channel.

    Args:
        image (np.ndarray): Input image. Can be gray-scale or in color.
        reference (np.ndarray): Image to match histogram of. Must have the same number of channels as image.
        channel_axis (int | None): If None, the image is assumed to be a grayscale (single channel) image.
            Otherwise, this indicates which axis of the array corresponds to channels.

    Returns:
        np.ndarray: Transformed input image.

    Raises:
        ValueError: Thrown when the number of channels in the input image and the reference differ.

    """
    if reference.dtype != np.uint8:
        reference = from_float(reference, np.uint8)

    matched = np.empty(image.shape, dtype=np.uint8)

    num_channels = image.shape[-1]

    for channel in range(num_channels):
        matched_channel = _match_cumulative_cdf(image[..., channel], reference[..., channel]).astype(np.uint8)
        matched[..., channel] = matched_channel

    return matched


def _match_cumulative_cdf(source: np.ndarray, template: np.ndarray) -> np.ndarray:
    src_lookup = source.reshape(-1)
    src_counts = np.bincount(src_lookup)
    tmpl_counts = np.bincount(template.reshape(-1))

    # omit values where the count was 0
    tmpl_values = np.nonzero(tmpl_counts)[0]
    tmpl_counts = tmpl_counts[tmpl_values]

    # calculate normalized quantiles for each array
    src_quantiles = np.cumsum(src_counts) / source.size
    tmpl_quantiles = np.cumsum(tmpl_counts) / template.size

    interp_a_values = np.interp(src_quantiles, tmpl_quantiles, tmpl_values)
    return interp_a_values[src_lookup].reshape(source.shape).astype(np.uint8)
