from __future__ import annotations

import frangi_3d as fg_
from glial_cmp import glial_cmp_t
import map_labeling as ml_
from type import array_t, site_h

from typing import Optional, Sequence, Tuple

import numpy as np_
import skimage.filters as fl_
import skimage.measure as ms_
import skimage.morphology as mp_
from scipy import ndimage as im_


min_area_c = 100


class extension_t(glial_cmp_t):
    #
    # bmp=boolean map
    # lmp=labeled map (intXX or uintXX array)
    # map=extension map (map=binary, int8 or uint8 array))
    # ep_=end point
    # soma_uid: connected to a soma somewhere upstream (as opposed to downstream extensions)
    # extensions: downstream (as opposed to being upstreamed connected)
    #
    __slots__ = ("scales", "end_points", "soma_uid", "__cache__")

    def __init__(self):
        #
        super().__init__()
        self.scales = None
        self.end_points = None
        self.soma_uid = None
        self.__cache__ = None

    @classmethod
    def FromMap(cls, lmp: array_t, scales: array_t, uid: int) -> extension_t:
        #
        instance = cls()

        bmp = lmp == uid
        instance.InitializeFromMap(bmp, uid)
        end_point_map = extension_t.EndPointMap(bmp)
        end_point_lmp = end_point_map * lmp
        instance.end_points = (end_point_lmp == uid).nonzero()
        instance.scales = scales[instance.sites]
        instance.__cache__ = {}

        return instance

    @property
    def is_unconnected(self) -> bool:
        #
        return self.soma_uid is None

    @property
    def end_points_as_array(self) -> array_t:
        #
        pty_name = "end_points_as_array"
        if pty_name not in self.__cache__:
            self.__cache__[pty_name] = np_.array(self.end_points)

        return self.__cache__[pty_name]

    def EndPointsForSoma(
        self, soma_uid: int, influence_map: array_t
    ) -> Tuple[site_h, ...]:
        #
        ep_bmp = influence_map[self.end_points] == soma_uid  # bmp=boolean map
        if ep_bmp.any():
            end_point_idc = ep_bmp.nonzero()[0]
            end_points = self.end_points_as_array[:, end_point_idc]

            return tuple(zip(*end_points.tolist()))

        return ()

    def BackReferenceSoma(self, glial_cmp: glial_cmp_t) -> None:
        #
        if isinstance(glial_cmp, extension_t):
            self.soma_uid = glial_cmp.soma_uid
        else:
            self.soma_uid = glial_cmp.uid

    def __str__(self) -> str:
        #
        if self.extensions is None:
            n_extensions = 0
        else:
            n_extensions = self.extensions.__len__()

        return (
            f"Ext.{self.uid}, "
            f"sites={self.sites[0].__len__()}, "
            f"endpoints={self.end_points[0].__len__()}, "
            f"soma={self.soma_uid}, "
            f"extensions={n_extensions}"
        )

    @staticmethod
    def ExtensionWithSite(
        extensions: Sequence[extension_t], site: site_h
    ) -> Optional[extension_t]:
        #
        for extension in extensions:
            if site in tuple(zip(*extension.sites)):
                return extension

        return None

    @staticmethod
    def EnhancedForDetection(
        image: array_t, in_parallel: bool = False
    ) -> Tuple[array_t, array_t]:
        #
        import os.path as ph_

        if ph_.exists("./frangi.npz"):
            print("/!\\ Reading from precomputed data file")
            loaded = np_.load("./frangi.npz")
            enhanced_img = loaded["enhanced_img"]
            scale_map = loaded["scale_map"]

            return enhanced_img, scale_map

        preprocessed_img = im_.morphology.white_tophat(
            image, size=2, mode="constant", cval=0.0, origin=0
        )

        enhanced_img, scale_map = fg_.FrangiEnhancement(
            preprocessed_img,
            scale_range=(0.1, 3),
            scale_step=1,
            alpha=0.8,
            beta=0.5,
            frangi_c=500,
            bright_on_dark=True,
            in_parallel=in_parallel,
        )

        np_.savez_compressed(
            "./frangi.npz", enhanced_img=enhanced_img, scale_map=scale_map
        )

        return enhanced_img, scale_map

    @staticmethod
    def CoarseMap(image: array_t, low: float, high: float) -> array_t:
        #
        result = __HysterisisImage__(image, low, high)
        result = __MorphologicalCleaning__(result)

        return result

    @staticmethod
    def FilteredCoarseMap(map_: array_t) -> array_t:
        #
        result = map_.copy()
        lmp = ms_.label(map_)

        for region in ms_.regionprops(lmp):
            if region.area <= min_area_c:
                region_sites = (lmp == region.label).nonzero()
                result[region_sites] = 0

        return result

    @staticmethod
    def FineMapFromCoarseMap(coarse_map: array_t) -> array_t:
        #
        # Might contain True-voxels that could be removed w/o breaking connectivity
        result = mp_.skeletonize_3d(coarse_map.astype(np_.uint8, copy=False))

        return result.astype(np_.int8, copy=False)

    @staticmethod
    def EndPointMap(map_: array_t) -> array_t:
        #
        part_map = ml_.PartLMap(map_)
        result = part_map == 1

        return result.astype(np_.int8)


def NormalizedImage(image: array_t) -> array_t:
    #
    middle_values = image[np_.logical_and(image > 0.0, image < image.max())]
    image_mean = middle_values.mean()
    result = image / image_mean

    return result


def __HysterisisImage__(image: array_t, low: float, high: float) -> array_t:
    #
    # low = 0.02
    # high = 0.04
    nonzero_sites = (image > 0).nonzero()
    nonzero_values = image[nonzero_sites]

    low = low * nonzero_values.min()
    high = high * image.max()
    # lowt = low*(max_image_f-min_image_f)+max_image_f
    # hight = high*(max_image_f- min_image_f)+min_image_f
    # lowt = (image_f >lowt).astype(int)
    # hight = (image_f <hight).astype(int)

    result = fl_.apply_hysteresis_threshold(image, low, high)
    result = result.astype(np_.int8, copy=False)

    return result


def __MorphologicalCleaning__(image: array_t) -> array_t:
    #
    result = image.copy()

    selem = mp_.disk(1)
    for dep in range(result.shape[0]):
        result[dep, :, :] = mp_.closing(result[dep, :, :], selem)
        result[dep, :, :] = mp_.opening(result[dep, :, :], selem)

    return result