import os

import numpy as np
import skimage.exposure
from ashlar.filepattern import FilePatternReader
from ashlar.reg import BioformatsMetadata, BioformatsReader
from skimage.filters import gaussian
from skimage.util import invert

from scportrait.processing.images._image_processing import rescale_image


class FilePatternReaderRescale(FilePatternReader):
    """Class for reading images based on a file pattern. If desired the images can be rescaled to a certain range while reading."""

    def __init__(
        self,
        path: str,
        pattern: str,
        overlap: float,
        pixel_size: int = 1,
        do_rescale: bool = False,
        no_rescale_channel: str = "Alexa488",
        rescale_range: tuple[int, int] = (1, 99),
        dtype=None,
    ):
        try:
            super().__init__(path, pattern, overlap, pixel_size=pixel_size)
        except (OSError, FileNotFoundError):
            print(
                f"Error: Could not read images with the given pattern {pattern}. Please check the path {path} and pattern."
            )
            found_files = os.listdir(path)
            print(
                f"At the provided location the the files follow the naming convention:{found_files[0:max(5, len(found_files))]} "
            )

        self.do_rescale = do_rescale
        self.no_rescale_channel = no_rescale_channel
        self.rescale_range = rescale_range
        self.dtype_image = dtype

    @staticmethod
    def rescale(
        img: np.ndarray, rescale_range: tuple[int, int] = (1, 99), cutoff_threshold: int | None = None
    ) -> np.ndarray:
        return rescale_image(img, rescale_range, cutoff_threshold=cutoff_threshold)

    @staticmethod
    # placeholer method kept for compatibility with old code
    # should be reimplemented in the future to allow for more flexible illumination correction
    def correct_illumination(
        img: np.ndarray,
        sigma: float = 30,
        double_correct: bool = False,
        rescale_range: tuple[int, int] = (1, 99),
        cutoff_threshold: int | None = None,
    ) -> np.ndarray:
        img = rescale_image(img, rescale_range, cutoff_threshold=cutoff_threshold, return_float=True)

        # calculate correction mask
        correction = gaussian(img, sigma)
        correction = invert(correction)
        correction = skimage.exposure.rescale_intensity(correction, out_range=(0, 1))

        correction_lows = np.where(img > 0.5, 0, img) * correction
        img_corrected = skimage.exposure.rescale_intensity(img + correction_lows, out_range=(0, 1))

        if double_correct:
            correction_mask_highs = invert(correction)
            correction_mask_highs_02 = skimage.exposure.rescale_intensity(
                np.where(img_corrected < 0.5, 0, img_corrected) * correction_mask_highs
            )
            img_corrected_double = skimage.exposure.rescale_intensity(img_corrected - 0.25 * correction_mask_highs_02)

            return (img_corrected_double * 65535).astype("uint16")
        else:
            return (img_corrected * 65535).astype("uint16")

    def read(self, series, c) -> np.ndarray:
        img = super().read(series, c)

        # check to ensure we are always using the same dtype
        if self.dtype_image is not None:
            if not isinstance(img.dtype, self.dtype_image):
                Warning(f"Found image with dtype {img.dtype}. Automatically converting to {self.dtype_image}")
                img = img.astype(self.dtype_image)
        else:
            # first time reading an image set dtype if none is supplied
            print("Setting dtype to", img.dtype)
            self.dtype_image = type(img.dtype)

        if not self.do_rescale:
            return img
        else:
            if c not in self.no_rescale_channel:
                # get rescale_range for channel c
                if isinstance(self.rescale_range, dict):
                    rescale_range = self.rescale_range[c]
                else:
                    rescale_range = self.rescale_range

                # actually read the image with rescaling applied
                return self.rescale(img, rescale_range=rescale_range)
            else:
                return img


class BioformatsMetadataRescale(BioformatsMetadata):
    """Reimplementation of BioformatsMetadata class to provide same parameters as contained in FilePatternReaderRescale"""

    def __init__(self, path):
        super().__init__(path)

    @property
    def channel_map(self):
        n_channels = self._metadata.getChannelCount(0)
        channel_names = []

        for id in range(n_channels):
            channel_names.append(self._metadata.getChannelName(0, id))
        channel_map = dict(zip(list(range(n_channels)), channel_names, strict=False))

        return channel_map


class BioformatsReaderRescale(BioformatsReader):
    """Class for reading images from Bioformats files (e.g. nd2).

    If desired the images can be rescaled to a certain range while reading.
    """

    def __init__(self, path, plate=None, well=None, do_rescale=False, no_rescale_channel=None, rescale_range=(1, 99)):
        self.path = path
        self.metadata = BioformatsMetadataRescale(self.path)
        self.metadata.set_active_plate_well(plate, well)
        self.do_rescale = do_rescale
        self.no_rescale_channel = no_rescale_channel
        self.rescale_range = rescale_range

    @staticmethod
    def rescale(img, rescale_range=(1, 99)):
        img = skimage.util.img_as_float32(img)
        cutoff1, cutoff2 = rescale_range

        if img.max() > (40000 / 65535):
            _img = img.copy()
            _img[_img > (10000 / 65535)] = 0
            p1 = np.percentile(_img, cutoff1)
            p99 = np.percentile(_img, cutoff2)
        else:
            p1 = np.percentile(img, cutoff1)
            p99 = np.percentile(img, cutoff2)

        img = skimage.exposure.rescale_intensity(img, in_range=(p1, p99), out_range=(0, 1))
        return (img * 65535).astype("uint16")

    def read(self, series, c):
        self.metadata._reader.setSeries(self.metadata.active_series[series])
        index = self.metadata._reader.getIndex(0, c, 0)
        byte_array = self.metadata._reader.openBytes(index)
        dtype = self.metadata.pixel_dtype  # no check required here as dtype has to be the same for all
        shape = self.metadata.tile_size(series)
        img = np.frombuffer(byte_array.tostring(), dtype=dtype).reshape(shape)

        if not self.do_rescale:
            return img
        else:
            if c not in self.no_rescale_channel:
                # get rescale_range for channel c
                if isinstance(self.rescale_range, dict):
                    rescale_range = self.rescale_range[c]
                else:
                    rescale_range = self.rescale_range

                # actually read the image with rescaling applied
                return self.rescale(img, rescale_range=rescale_range)
            else:
                return img
