import abc
from pathlib import Path
from typing import Iterable

import cv2
import numpy as np
from csbdeep.utils import normalize

from neuralib.argp import AbstractParser, argument

__all__ = ['AbstractSegmentationOption']


class AbstractSegmentationOption(AbstractParser, metaclass=abc.ABCMeta):
    DESCRIPTION = 'ABC for cellular segmentation'

    file: Path = argument('-F', '--file', help='image file path')
    directory: Path = argument('-D', '--dir', help='images directory for batch processing')
    directory_suffix: str = argument('--suffix', default='.tif',
                                     choices=['.tif', '.tiff', '.png'],
                                     help='suffix in batch mode')
    save_ij_roi: bool = argument('--ij-roi', '--roi', help='if save also the imageJ/Fiji compatible .roi file')
    force_re_eval: bool = argument('--force-eval', '--re', help='force re-evaluate the result')

    model: str = argument('-M', '--model', metavar='MODEL', help='which pretrained model')
    no_normalize: bool = argument('--no-norm', help='NOT DO Percentile-based image normalization for eval')
    napari_view: bool = argument('--napari', help='view in napari')

    def post_parsing(self):
        """check args is valid"""
        if not self.file_mode and not self.batch_mode:
            raise ValueError('Either use "-F in file" or "-D in directory" for processing')

        if self.file and self.directory:
            raise ValueError('Only "-F in file" or "-D in directory" should be specified, not both')

        if self.directory and self.napari_view:
            raise ValueError('napari view only used in single file mode')

    @property
    def file_mode(self) -> bool:
        """Flag file mode"""
        return self.file is not None

    @property
    def batch_mode(self) -> bool:
        """Flag batch mode"""
        return self.directory is not None

    @staticmethod
    def _as_grayscale(file: str) -> np.ndarray:
        img = cv2.imread(file)
        if img.ndim == 3:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return img

    def raw_image(self) -> np.ndarray:
        """Load image from file and convert to grayscale

        :return: `Array[float, [H, W]]`
        """
        if not self.file.is_file():
            raise ValueError('not a file')

        return self._as_grayscale(str(self.file))

    def normalize_image(self) -> np.ndarray:
        """
        Normalize the image

        :return: `Array[float, [H, W]]`
        """
        img = self.raw_image()

        return normalize(img, clip=True)

    def foreach_raw_image(self) -> Iterable[tuple[Path, np.ndarray]]:
        """
        Load image from a directory and convert to grayscale

        :return: Tuple of filepath and image `Array[float, [H, W]]`
        """
        for file in self.directory.glob(f'*{self.directory_suffix}'):
            img = self._as_grayscale(str(file))
            yield file, img

    def foreach_normalize_image(self) -> Iterable[tuple[Path, np.ndarray]]:
        """
        Normalize the image in batch mode

        :return: Tuple of filepath and image `Array[float, [H, W]]`
        """
        for name, raw in self.foreach_raw_image():
            yield name, normalize(raw, clip=True)

    @abc.abstractmethod
    def seg_output(self, filepath: Path) -> Path:
        """
        Get segmented output save path

        :param filepath: filepath for image
        :return: segmented output save path
        """
        pass

    def ij_roi_output(self, filepath: Path) -> Path:
        """
        Get imageJ/Fiji ``.roi`` output save path

        :param filepath: filepath for image
        :return: ij roi output save path
        """
        return filepath.with_suffix('.roi')

    @abc.abstractmethod
    def eval(self) -> None:
        """eval the model in single file or batch files, and save the results"""
        pass

    @abc.abstractmethod
    def launch_napari(self, **kwargs):
        """napari viewer"""
        pass
