# Copyright CNRS/Inria/UNS
# Contributor(s): Eric Debreuve (since 2019), Morgane Nadal (2020)
#
# 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 re as re_
import numpy as np_
from copy import deepcopy

from brick.component.soma import soma_t
from sklgraph.skl_graph import skl_graph_t
from typing import Tuple, Dict, List, Union


def FindGraphsRootWithNodes(somas: Tuple[soma_t, ...], ext_skl_graph: skl_graph_t) -> skl_graph_t:
    """
    Find the roots of the {extension+connexion} graphs to be lined to the soma.
    Add a key "root" (bool) in the dict of nodes attributes.
    """

    node_degree_bool = tuple(degree == 1 for _, degree in ext_skl_graph.degree)
    node_coord = tuple(xyz for xyz, _ in ext_skl_graph.degree)

    # get the coordinates of the nodes (x,y,z)
    coordinates = GetNodesCoordinates(node_coord)

    # get a list with elements = (soma_uid, extension_uid, root coordinates)
    roots = GetListRoots(somas)

    # for each node in the graph, search among the degree 1 nodes the nodes that are roots (linked to soma)
    for node in range(len(coordinates)):
        if node_degree_bool[node]:
            # compare the coor with end points
            for ext_root in roots:
                if ext_root[2] == coordinates[node]:
                    # add the attribute 'root' = True in the dict of nodes attributes
                    ext_skl_graph._node[node_coord[node]]['root'] = True
                    # add the id of soma and of the extension in the nodes attributes dict
                    ext_skl_graph._node[node_coord[node]]['soma_uid'] = ext_root[0]
                    ext_skl_graph._node[node_coord[node]]['ext_uid'] = ext_root[1]
                else:
                    try:
                        # verify if the node has already been covered to avoid overwriting on a True value:
                        if ext_skl_graph._node[node_coord[node]]['root'] is not True:
                            ext_skl_graph._node[node_coord[node]]['root'] = False
                    except:
                        # add the attribute 'root' = False for non-root endpoints in the dict of nodes attributes
                        ext_skl_graph._node[node_coord[node]]['root'] = False
        else:
            # all the nodes with degree > 1 are not roots
            ext_skl_graph._node[node_coord[node]]['root'] = False

    return ext_skl_graph

    # TODO: add to every edge the id of the extension


def FindGraphsRootWithEdges(somas: Tuple[soma_t,...], ext_skl_graph: skl_graph_t, ext_nfo) -> list:
    """
    Finds the soma roots of the graph extension.
    """
    # For a given soma, find the roots of the graphs
    list_root_nodes = []

    for soma in somas:

        dict_root_nodes = {}
        root_nodes = {}

        # Finds the primary extensions
        primary_extension_uids = tuple(extension.uid for extension in soma.extensions)

        # List of the degree 1 nodes of the graph
        for node1_id, node2_id, edge_nfo in ext_skl_graph.edges.data('as_edge_t'):
            if (ext_skl_graph.degree[node1_id] == 1) or (ext_skl_graph.degree[node2_id] == 1):

                # Find the pixels of the terminal extension
                sites = ext_nfo['lmp'][edge_nfo.sites]
                ext_uid = np_.unique(sites[np_.where(sites > 0)])
                # np_.where(sites>0) because ext_nfo['lmp'] do not contain the connexions

                # Add the ext uid to the extension -- use deepcopy to avoid pb with the iteration on the graph ?
                # ext_skl_graph.get_edge_data(node1_id, node2_id)['Ext. uid'] = ext_uid

                if ext_uid in primary_extension_uids:
                    if ext_skl_graph.degree[node1_id] == 1:
                        root_node = node1_id
                    else:
                        root_node = node2_id

                    # Search for the sites in the image of the soma extension
                    # edge_sites = set(zip(*edge_nfo.sites))

                    # for extension in soma.extensions:   # In this case, intersections is never empty : we must take the
                    # contour sites of the soma, extended to the neighbor voxels
                    #    if extension.uid == ext_uid:
                    #       soma_sites = set(zip(*extension.sites))

                    root_node_coor = GetNodesCoordinates(tuple([root_node]))[0]

                    root_sites = tuple(
                        (root_node_coor[0] + i, root_node_coor[1] + j, root_node_coor[2] + k)
                        for i in (-1, 0, 1)
                        for j in (-1, 0, 1)
                        for k in (-1, 0, 1)
                        if i != 0 or j != 0 or k != 0)

                    soma_contour_sites = set(soma.contour_points)
                    # TODO: Which is better ? Extend the potential root node at every iteration or to extend once the
                    #  soma contour ?
                    intersections = soma_contour_sites.intersection(root_sites)

                    # if the graph root sites are included in the soma extensions sites
                    if len(intersections) > 0:
                        # Keep the info of the root node
                        root_nodes['Ext. ' + str(ext_uid[0])] = root_node  # By construction, only one root node possible for an ext
        dict_root_nodes["Soma " + str(soma.uid)] = root_nodes
        list_root_nodes.append(dict_root_nodes)

    return list_root_nodes


def GetNodesCoordinates(node_coord: Tuple[str, ...]) -> list:
    """
    Input: nodes attributes -> Tuple('x1-y1-z1', 'x2-y2-z2', ...) .
    Output: coordinates -> List[Tuple(x1,y1,z1), Tuple(x2,y2,z2), ...]
    """
    coord = []
    for c in node_coord:
        coord.append(c)

    for node in range(len(node_coord)):
        coord_node = coord[node]
        pattern = '\d+'
        coord_node = re_.findall(pattern, coord_node)
        coor = []
        for i in range(3):
            coor.append(int(coord_node[i]))
        coor = tuple(coor)
        coord[node] = coor

    return coord


def GetListRoots(somas: soma_t) -> list:
    """
    Gives a list containing the following information for all somas: [soma id: int, extension id: int, root = (x,y,z): tuple]
    """
    roots = []
    for soma in somas:
        for ext_id, ext_root in enumerate(soma.ext_root):
            roots.append((soma.uid, soma.extensions[ext_id].uid, ext_root[0]))
    return roots