import dataclasses
from typing import List, Optional

import cv2
import numpy as np


@dataclasses.dataclass
class OtsuThreshold(object):
    """Otsu thresholding.
    min_value: the value of class1 pixels after masking.
    max_value: the value of class2 pixels after masking.
    """

    min_value: int = 0
    max_value: int = 255

    def _convert_gray_image(self, image: np.ndarray) -> np.ndarray:
        if image.ndim == 2:
            # already gray image
            return image
        else:
            # convert color image to gray image
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            return gray

    def _check_mask(self, mask: np.ndarray) -> np.ndarray:
        if mask.ndim != 2:
            raise ValueError("Invalid mask.")

        unique_value = np.unique(mask)
        if len(unique_value) != 2:
            raise ValueError(
                "Mask should be binarized and have two values: [0, 1] or [0, 255]"
            )

        if np.all(unique_value == np.array([0, 1])):
            mask *= 255
            return mask
        elif np.all(unique_value == np.array([0, 255])):
            return mask
        else:
            raise ValueError("Values of the mask should be [0, 1] or [0, 255].")

    def _mask_gray_image(
        self,
        gray: np.ndarray,
        mask: np.ndarray,
    ) -> np.ndarray:
        """Masking gray image
        ignoring the pixels whose value in mask are 255.

        output will be 1D array
        """
        masked_gray = gray[np.where(mask == 255)]
        return masked_gray

    def _calc_hist(self, gray: np.ndarray) -> List[int]:
        hist = [np.sum(gray == i) for i in range(256)]
        return hist

    def _find_max_separation_th(
        self,
        hist: List[int],
    ) -> int:
        """Brute-force searching for the maximum separation
        The degree of separation is defined as;
        the degree of separation = inter-class variance / intra-class variance

        all data variance = inter-class vjariance + intra-class variance
        then, the degree of separation
        = inter-class variance / (all data variance - inter-class variance)

        to maximize the degree of separation, maximize the inter-class variance
        inter-class variance is calculated as below.
        inter-class variance = n1 * n2 * (mu1- mu2) ** 2

        Herein, we use inter-class variance as `separation`
        """
        max_separation = 0.0
        max_separation_th = 0

        for th in range(256):
            # the number of pixels of class1 and class2
            n1 = sum(hist[:th])
            n2 = sum(hist[th:])

            # calculating the average values of pixels each class.
            if n1 == 0:
                mean1 = 0.0
            else:
                mean1 = sum([i * hist[i] for i in range(th)]) / n1

            if n2 == 0:
                mean2 = 0.0
            else:
                mean2 = sum([i * hist[i] for i in range(th, 256)]) / n2

            separation = n1 * n2 * (mean1 - mean2) ** 2

            # record maximum separation and the threshold
            if separation > max_separation:
                max_separation = separation
                max_separation_th = th

        return max_separation_th

    def _binaraize(self, gray: np.ndarray, th: int) -> np.ndarray:
        gray[gray < th] = self.min_value
        gray[gray >= th] = self.max_value
        return gray

    def process(
        self,
        image: np.ndarray,
        mask: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """Processing Otsu thresholding
        Args:
            image: RGB image array (H, W, 3) or gray-scale image array (H, W)
            mask: the mask array with the same size as the image. (H, W)
                The pixel values should be in {0, 1} or {0, 255}.
                The pixel with 0 value will be ignored
                when the histogram and the degree of separation is calculated.
        Return:
            mask: generated by otsu thresholding.
        """
        gray = self._convert_gray_image(image)

        if mask is not None:
            mask = self._check_mask(mask)
            masked_gray = self._mask_gray_image(gray, mask)
        else:
            masked_gray = gray.copy()

        # calculate histgram
        hist = self._calc_hist(masked_gray)

        # find the threshold which maximizes the separation
        max_separation_th = self._find_max_separation_th(hist)

        # binarize
        gray = self._binaraize(gray, max_separation_th)

        return gray


if __name__ == "__main__":
    processor = OtsuThreshold()

    img = cv2.imread("../imgs/2007_000032.jpg")
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    res = processor.process(gray, mask=None)
    cv2.imshow("window1", res)
    cv2.waitKey(-1)

    # masked otsu
    mask = np.zeros_like(gray, dtype=np.uint8)
    mask[0:200, -200:] = 1

    res = processor.process(gray, mask=mask)

    cv2.imshow("window2", res)
    cv2.waitKey(-1)
