from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING: from phenotypic import Image

import numpy as np
from skimage.morphology import disk, square, white_tophat, cube, ball, diamond

from phenotypic.abc_ import ImageEnhancer


class WhiteTophatEnhancer(ImageEnhancer):
    def __init__(self, shape: str = 'diamond', radius: int = None):
        self.shape = shape
        self.radius = radius

    def _operate(self, image: Image) -> Image:
        white_tophat_results = white_tophat(
                image.enh_gray[:],
                footprint=self._get_footprint(
                        self._get_footprint_radius(detection_matrix=image.enh_gray[:]),
                ),
        )
        image.enh_gray[:] = image.enh_gray[:] - white_tophat_results

        return image

    def _get_footprint_radius(self, detection_matrix: np.ndarray) -> int:
        if self.radius is None:
            return int(np.min(detection_matrix.shape)*0.004)
        else:
            return self.radius

    def _get_footprint(self, radius: int) -> np.ndarray:
        match self.shape:
            case 'disk':
                return disk(radius=radius)
            case 'square':
                return square(width=radius*2)
            case 'diamond':
                return diamond(radius=radius)
            case 'sphere':
                return ball(radius)
            case 'cube':
                return cube(radius*2)
