from PIL import Image, ImageChops
from dataclasses import dataclass
from typing import TYPE_CHECKING
import numpy as np
import scipy.ndimage
from skimage.metrics import structural_similarity

if TYPE_CHECKING:
    from .patch_editor import EditorCtx


@dataclass
class DiffAlgorithm:
    """
    Base class for difference calculation algorithms.

    Subclasses should implement the `calculate` method to return a heatmap
    of differences between two images.
    """

    ctx: "EditorCtx"
    half: int = 0
    full: int = 255

    def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
        raise NotImplementedError

    def get_threshold(self, tolerance: float) -> int:
        if tolerance <= 1.0:
            return int(self.half * tolerance)
        else:
            return int(self.half + (self.full - self.half) * (tolerance - 1.0))

    def apply_alpha_mask(self, arr: np.ndarray, img2: Image.Image) -> np.ndarray:
        if img2.mode == 'RGBA':
            alpha_mask = np.array(img2.split()[-1]) > 0
            arr *= alpha_mask
        return arr


@dataclass
class HeatmapDifference(DiffAlgorithm):
    """
    Calculates the difference between two images using ImageChops.difference.
    This is a simple and fast algorithm that computes the absolute difference
    for each pixel.
    """

    half: int = 10
    full: int = 40

    def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
        diff = ImageChops.difference(img1.convert("RGB"), img2.convert("RGB"))
        diff_l = diff.convert('L')
        arr = np.array(diff_l, dtype=np.float32)
        arr = self.apply_alpha_mask(arr, img2)

        if self.ctx.region_size > 1:
            arr = scipy.ndimage.uniform_filter(arr, size=self.ctx.region_size, mode='reflect')

        return arr


@dataclass
class HeatmapL2Norm(DiffAlgorithm):
    """
    Calculates the L2 norm (Euclidean distance) between the pixel values of two images.
    This method is more sensitive to color variations than simple difference.
    """

    half: int = 10
    full: int = 100

    def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
        a1 = np.array(img1.convert('RGB'), dtype=np.float32)
        a2 = np.array(img2.convert('RGB'), dtype=np.float32)

        diff_sq = np.sum((a1 - a2) ** 2, axis=2)
        diff_sq = self.apply_alpha_mask(diff_sq, img2)

        if self.ctx.region_size > 1:
            mse = scipy.ndimage.uniform_filter(diff_sq, size=self.ctx.region_size, mode='reflect')
        else:
            mse = diff_sq

        return np.sqrt(mse)


@dataclass
class HeatmapSSIM(DiffAlgorithm):
    """
    Calculates the Structural Similarity Index (SSIM) between two images.
    SSIM is a perceptual metric that quantifies the similarity between two images,
    considering luminance, contrast, and structure.
    """

    half: int = 40
    full: int = 100

    def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
        a1 = np.array(img1.convert('RGB'))
        a2 = np.array(img2.convert('RGB'))

        win_size = self.ctx.region_size if self.ctx.region_size % 2 == 1 else self.ctx.region_size + 1
        if win_size < 3:
            win_size = 3

        kwargs = {'channel_axis': 2, 'data_range': 255, 'win_size': win_size}

        _, diff = structural_similarity(a1, a2, full=True, **kwargs)

        diff_map = (1 - diff) * 255

        if diff_map.ndim == 3:
            diff_map = np.mean(diff_map, axis=2)

        diff_map = self.apply_alpha_mask(diff_map, img2)
        return diff_map