soma.py

from __future__ import annotations

from glial_cmp import glial_cmp_t
import map_labeling as ml_
from type import array_t, py_array_picker_h, site_h

from collections import namedtuple as namedtuple_t
import sys as sy_
from typing import Optional, Sequence, Tuple

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


som_ext_path_t = namedtuple_t("som_ext_path_t", "extension length path idx")


class soma_t(glial_cmp_t):
    #
    # bmp=boolean map
    # lmp=labeled map (intXX or uintXX array)
    # map=extension map (map=binary, int8 or uint8 array))
    #
    __slots__ = ("contour_points",)

    def __init__(self):
        #
        super().__init__()
        self.contour_points = None

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

        bmp = lmp == uid
        instance.InitializeFromMap(bmp, uid)
        contour_map = soma_t.ContourMap(bmp)
        contour_lmp = contour_map * lmp
        instance.contour_points = tuple(zip(*(contour_lmp == uid).nonzero()))

        return instance

    def Extensions(self, max_level: int = sy_.maxsize) -> Tuple[glial_cmp_t]:
        #
        # max_level=1: directly connected extensions
        # ...
        soma_ext = self.extensions.copy()

        current_level_idx = 0
        for level in range(2, max_level + 1):
            next_level_idx = soma_ext.__len__()
            if next_level_idx == current_level_idx:
                break

            for extension in soma_ext[current_level_idx:]:
                soma_ext.extend(extension.extensions)

            current_level_idx = next_level_idx

        return tuple(soma_ext)

    def ContourPointsCloseTo(
        self, point: site_h, max_distance: float
    ) -> Tuple[Optional[Tuple[site_h, ...]], Optional[py_array_picker_h]]:
        #
        points = tuple(
            contour_point
            for contour_point in self.contour_points
            if (np_.subtract(point, contour_point) ** 2).sum() <= max_distance
        )

        if points.__len__() > 0:
            points_as_picker = tuple(zip(*points))
        else:
            points = None
            points_as_picker = None

        return points, points_as_picker

    def ExtensionPointsCloseTo(
        self, point: site_h, max_distance: float
    ) -> Tuple[Optional[Tuple[site_h, ...]], Optional[py_array_picker_h]]:
        #
        # Leave connection paths apart because only detected extension pieces
        # (as opposed to invented=connection paths) are considered valid connection targets
        #
        points = []
        for extension in self.Extensions():
            ext_sites = tuple(zip(*extension.sites))
            points.extend(
                ext_point
                for ext_point in ext_sites
                if (np_.subtract(point, ext_point) ** 2).sum() <= max_distance
            )
        points = tuple(points)

        if points.__len__() > 0:
            points_as_picker = tuple(zip(*points))
        else:
            points = None
            points_as_picker = None

        return points, points_as_picker

    # def Extend(
    #     self, extensions: Sequence[extension_t], dist_to_soma: array_t, costs: array_t
    # ) -> None:
    #     #
    #     candidate_ext_nfo = []  # eps=end points
    #     for extension in extensions:
    #         new_candidates = extension.EndPointsForSoma(self.uid)
    #         candidate_ext_nfo.extend(
    #             (end_point, extension) for end_point in new_candidates
    #         )
    #     candidate_ext_nfo.sort(key=lambda elm: dist_to_soma[elm[0]])
    #     unconnected_candidates = list(
    #         filter(lambda elm: elm[1].soma_uid is None, candidate_ext_nfo)
    #     )
    #
    #     should_look_for_connections = True
    #     while should_look_for_connections:
    #         som_ext_path_nfos = []
    #         for ep_idx, (ext_end_point, extension) in enumerate(unconnected_candidates):
    #             path, length = self.ShortestPathTo(
    #                 extension.uid, ext_end_point, ep_idx, costs
    #             )
    #             if length <= max_weighted_dist:
    #                 som_ext_path_nfos.append(
    #                     som_ext_path_t(
    #                         extension=extension, length=length, path=path, idx=ep_idx
    #                     )
    #                 )
    #                 if path.__len__() == 2:
    #                     break
    #
    #         if som_ext_path_nfos.__len__() > 0:
    #             som_ext_path_nfos.sort(key=lambda nfo: nfo.length)
    #             shorest_path = som_ext_path_nfos[0]
    #             self.ExtendWith(shorest_path.extension, shorest_path.path, costs)
    #             del unconnected_candidates[shorest_path.idx]
    #         else:
    #             should_look_for_connections = False

    def BackReferenceSoma(self, glial_cmp: glial_cmp_t) -> None:
        #
        raise NotImplementedError("Soma do not need to back-reference a soma")

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

        return (
            f"Soma.{self.uid}, "
            f"area={self.sites[0].__len__()}, "
            f"contour points={self.contour_points.__len__()}, "
            f"extensions={n_extensions}"
        )

    @staticmethod
    def Map(image: array_t, low: float, high: float, selem: array_t) -> array_t:
        #
        # low = 10 #0.15
        # high = 67.4 # 0.7126
        #
        max_image = image.max()
        nonzero_sites = image.nonzero()
        nonzero_values = image[nonzero_sites]
        min_image = nonzero_values.min()

        low = low * (max_image - min_image) + min_image
        high = high * (max_image - min_image) + min_image
        result = fl_.apply_hysteresis_threshold(image, low, high)
        result = result.astype(np_.int8, copy=False)

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

        return result

    @staticmethod
    def FilteredMap(map_: array_t, min_area: int) -> array_t:

        result = map_.copy()
        lmp = ms_.label(map_)

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

        return result

    @staticmethod
    def ContourMap(map_: array_t) -> array_t:
        #
        part_map = ml_.PartLMap(map_)
        # Works because the background is labeled with 27
        result = part_map < 26

        return result.astype(np_.int8)

    @staticmethod
    def InfluenceMaps(map_: array_t) -> Tuple[array_t, array_t]:

        background = (map_ == 0).astype(np_.int8)
        dist_map, idx_map = im_.morphology.distance_transform_edt(
            background, return_indices=True
        )

        # obj_map = np_.empty_like(map_)
        # for row in range(0, obj_map.shape[0]):
        #     for col in range(0, obj_map.shape[1]):
        #         for dep in range(0, obj_map.shape[2]):
        #             obj_map[row, col, dep] = map_[
        #                 idx_map[0, row, col, dep],
        #                 idx_map[1, row, col, dep],
        #                 idx_map[2, row, col, dep],
        #             ]

        return dist_map, np_.array(map_[tuple(idx_map)])

    @staticmethod
    def SomasLMap(somas: Sequence[soma_t], with_extensions: bool = True) -> array_t:
        #
        shape = somas[0].img_shape
        result = np_.zeros(shape, dtype=np_.int64)

        for soma in somas:
            result[soma.sites] = soma.uid

            if with_extensions:
                for connection_path in filter(
                    lambda path: path is not None, soma.connection_path.values()
                ):
                    result[connection_path] = soma.uid
                for extension in soma.Extensions():
                    for connection_path in filter(
                        lambda path: path is not None,
                        extension.connection_path.values(),
                    ):
                        result[connection_path] = soma.uid
                    result[extension.sites] = soma.uid

        return result


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

    return result