from __future__ import annotations

from type import array_t, np_array_picker_h, py_array_picker_h, site_path_h

from typing import Dict, List, Optional, Tuple

import numpy as np_


class glial_cmp_t:
    #
    __slots__ = ("uid", "sites", "connection_path", "extensions", "img_shape")

    def __init__(self):
        #
        self.uid = None  # type: Optional[int]
        self.sites = None  # type: Optional[np_array_picker_h]
        self.connection_path = None  # type: Optional[Dict[int, py_array_picker_h]]
        self.extensions = None  # type: Optional[List[glial_cmp_t]]
        self.img_shape = None  # type: Optional[Tuple[int, ...]]

    def InitializeFromMap(self, bmp: array_t, uid: int) -> None:
        #
        self.uid = uid
        # sites: might contain voxels that could be removed w/o breaking connectivity
        self.sites = bmp.nonzero()
        self.connection_path = {}
        self.extensions = []
        self.img_shape = bmp.shape

    @property
    def has_extensions(self) -> bool:
        #
        return self.extensions.__len__() > 0

    def ExtendWith(
        self, extension: glial_cmp_t, through: site_path_h, costs: array_t
    ) -> None:
        #
        connection_path = tuple(zip(*through[1:-1]))
        if connection_path.__len__() == 0:
            connection_path = None
        extension_path = extension.sites

        self.connection_path[extension.uid] = connection_path
        self.extensions.append(extension)

        # TODO: put this outside
        extension.BackReferenceSoma(self)

        # TODO: Ideally, these paths should be dilated + put this outside
        # but in ext-ext connections, there must not be dilation around the current ext
        # (current ext plays the role of a soma in soma-ext step)
        if connection_path is not None:
            costs[connection_path] = np_.inf
        costs[extension_path] = np_.inf

    def BackReferenceSoma(self, glial_cmp: glial_cmp_t) -> None:
        #
        raise NotImplementedError