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:
    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))


@dataclass
class HeatmapDifference(DiffAlgorithm):
    half: int = 10
    full: int = 40

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

        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):
    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)

        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):
    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)

        return diff_map