"""Preprocessing methods"""
import logging
from typing import List, Tuple

import numpy as np
from PIL import Image, ImageFilter


from config import COLOR_RGB
# from enhance_config import ENHANCE_SETTINGS

LOGGING = logging.getLogger(__name__)


def preprocess_seg_mask(canvas_seg, real_seg: Image.Image = None) -> Tuple[np.ndarray, np.ndarray]:
    """Preprocess the segmentation mask.
    Args:
        canvas_seg: segmentation canvas
        real_seg (Image.Image, optional): segmentation mask. Defaults to None.
    Returns:
        Tuple[np.ndarray, np.ndarray]: segmentation mask, segmentation mask with overlay
    """
    # get unique colors in the segmentation
    image_seg = canvas_seg.image_data.copy()[:, :, :3]

    # average the colors of the segmentation masks
    average_color = np.mean(image_seg, axis=(2))
    mask = average_color[:, :] > 0
    if mask.sum() > 0:
        mask = mask * 1

    unique_colors = np.unique(image_seg.reshape(-1, image_seg.shape[-1]), axis=0)
    unique_colors = [tuple(color) for color in unique_colors]

    unique_colors = [color for color in unique_colors if np.sum(
        np.all(image_seg == color, axis=-1)) > 100]

    unique_colors_exact = [color for color in unique_colors if color in COLOR_RGB]

    if real_seg is not None:
        overlay_seg = np.array(real_seg)

        unique_colors = np.unique(overlay_seg.reshape(-1, overlay_seg.shape[-1]), axis=0)
        unique_colors = [tuple(color) for color in unique_colors]

        for color in unique_colors_exact:
            if color != (255, 255, 255) and color != (0, 0, 0):
                overlay_seg[np.all(image_seg == color, axis=-1)] = color
        image_seg = overlay_seg

    return mask, image_seg