# Copyright CNRS/Inria/UNS
# Contributor(s): Eric Debreuve (since 2018)
#
# eric.debreuve@cnrs.fr
#
# This software is governed by the CeCILL  license under French law and
# abiding by the rules of distribution of free software.  You can  use,
# modify and/ or redistribute the software under the terms of the CeCILL
# license as circulated by CEA, CNRS and INRIA at the following URL
# "http://www.cecill.info".
#
# As a counterpart to the access to the source code and  rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty  and the software's author,  the holder of the
# economic rights,  and the successive licensors  have only  limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using,  modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean  that it is complicated to manipulate,  and  that  also
# therefore means  that it is reserved for developers  and  experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and,  more generally, to use and operate it in the
# same conditions as regards security.
#
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.

import sklgraph.brick.edge as dg_
import sklgraph.brick.node as nd_
import sklgraph.skl_map as sm_

from typing import List, Tuple

import numpy as np_


array_t = np_.ndarray


def AssignNodeIDsToEdges(
    e_nodes: List[nd_.end_node_t],
    e_node_lmap: array_t,
    b_nodes: List[nd_.branch_node_t],
    b_node_lmap: array_t,
    edges: List[dg_.edge_t],
    edge_lmap: array_t,
) -> None:
    #
    # Not uint to allow for subtraction
    edge_parts = sm_.SkeletonPartMap(
        (edge_lmap > 0).astype(np_.int8), check_validity=None
    )

    # ep=edge end point; Leave < 2 since ==0 (length-1 edges) and ==1 (other edges) are needed
    for ep_coords in zip(*(edge_parts < 2).nonzero()):
        edge_label = edge_lmap[ep_coords]
        edge = edges[edge_label - 1]
        e_node_label = e_node_lmap[ep_coords]

        if e_node_label > 0:
            # End node-to-X edge (i.e., edge end point is also an end node)
            _AppendNodeIDToEdge(
                edge,
                np_.zeros(ep_coords.__len__(), dtype=np_.int64),
                e_nodes[e_node_label - 1],
            )
            if edge.sites[0].__len__() == 1:
                # End node-to-branch node edge (and there is a unique non-zero value in b_neighborhood)
                nh_slices_starts, b_neighborhood = _LMapNeighborhood(
                    b_node_lmap, ep_coords
                )
                b_node_label = b_neighborhood.max()
                b_coords = np_.transpose((b_neighborhood == b_node_label).nonzero())[0]
                _AppendNodeIDToEdge(
                    edge, nh_slices_starts.__add__(b_coords), b_nodes[b_node_label - 1]
                )
        else:
            nh_slices_starts, b_neighborhood = _LMapNeighborhood(b_node_lmap, ep_coords)
            force_after = False
            # Looping only for length-1, b-to-b edges
            for b_coords in zip(*b_neighborhood.nonzero()):
                b_node_label = b_neighborhood[b_coords]
                _AppendNodeIDToEdge(
                    edge,
                    nh_slices_starts.__add__(b_coords),
                    b_nodes[b_node_label - 1],
                    force_after=force_after,
                )
                force_after = not force_after


def _AppendNodeIDToEdge(
    edge: dg_.edge_t, b_coords: array_t, node: nd_.node_t, force_after: bool = False
) -> None:
    #
    edge.node_uids.append(node.uid)

    space_dim = b_coords.size
    first_site = tuple(edge.sites[idx_][0] for idx_ in range(space_dim))

    if isinstance(node, nd_.branch_node_t):
        sq_distance = (b_coords.__sub__(first_site) ** 2).sum()

        if sq_distance <= space_dim:
            edge.origin_node = node.uid

        if edge.sites[0].__len__() > 1:
            if sq_distance <= space_dim:
                edge.sites = tuple(
                    np_.hstack((b_coords[idx_], edge.sites[idx_]))
                    for idx_ in range(space_dim)
                )
            else:
                edge.sites = tuple(
                    np_.hstack((edge.sites[idx_], b_coords[idx_]))
                    for idx_ in range(space_dim)
                )
        elif force_after:
            edge.sites = tuple(
                np_.hstack((edge.sites[idx_], b_coords[idx_]))
                for idx_ in range(space_dim)
            )
        else:
            edge.sites = tuple(
                np_.hstack((b_coords[idx_], edge.sites[idx_]))
                for idx_ in range(space_dim)
            )
    #
    elif np_.array_equal(first_site, node.position):
        edge.origin_node = node.uid


def _LMapNeighborhood(lmsk: array_t, site: Tuple[int, ...]) -> Tuple[array_t, array_t]:
    #
    slices_starts = tuple(max(site[idx_] - 1, 0) for idx_ in range(site.__len__()))
    slices = tuple(
        slice(slices_starts[idx_], min(site[idx_] + 2, lmsk.shape[idx_]))
        for idx_ in range(site.__len__())
    )
    neighborhood = lmsk[slices]

    return np_.array(slices_starts, dtype=np_.int64), neighborhood