#!/usr/bin/env python
#-*- coding:utf-8 -*-
#
# This file is part of the NNGT project to generate and analyze
# neuronal networks and their activity.
# Copyright (C) 2015-2017  Tanguy Fardet
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

""" iGraph subclassing """

from collections import OrderedDict

import numpy as np
import scipy.sparse as ssp

import nngt
from nngt.lib import InvalidArgument, nonstring_container, BWEIGHT
from .base_graph import BaseGraph, BaseProperty



#-----------------------------------------------------------------------------#
# Properties
#------------------------
#

class _IgNProperty(BaseProperty):

    '''
    Class for generic interactions with nodes properties (igraph)
    '''

    def __getitem__(self, name):
        if super(_IgNProperty, self).__getitem__(name) in ('string', 'object'):
            return np.array(np.array(self.parent().vs[name]), dtype=object)
        else:
            return np.array(self.parent().vs[name])

    def __setitem__(self, name, value):
        size = self.parent().vcount()
        if name in self:
            if len(value) == size:
                self.parent().vs[name] = value
            else:
                raise ValueError("A list or a np.array with one entry per "
                                 "node in the graph is required")
        else:
            raise InvalidArgument("Attribute does not exist yet, use "
                                  "set_attribute to create it.")

    def new_attribute(self, name, value_type, values=None, val=None):
        if val is None:
            if value_type == "int":
                val = int(0)
            elif value_type == "double":
                val = 0.
            elif value_type == "string":
                val = ""
            else:
                val = None
                value_type = "object"
        if values is None:
            values = [val for _ in range(self.parent().vcount())]
        # store name and value type in the dict
        super(_IgNProperty,self).__setitem__(name, value_type)
        # store the real values in the attribute
        self[name] = values
        self._num_values_set[name] = len(values)

    def set_attribute(self, name, values, nodes=None):
        '''
        Set the node attribute.
        
        Parameters
        ----------
        name : str
            Name of the node attribute.
        values : array, size N
            Values that should be set.
        nodes : array-like, optional (default: all nodes)
            Nodes for which the value of the property should be set. If `nodes`
            is not None, it must be an array of size N.
        '''
        num_nodes = self.parent().vcount()
        num_n = len(nodes) if nodes is not None else num_nodes
        if num_n == num_nodes:
            self[name] = values
        else:
            if num_n != len(values):
                raise ValueError("`nodes` and `nodes` must have the same "
                                 "size; got respectively " + str(num_n) + \
                                 " and " + str(len(values)) + "entries.")
            if self._num_values_set[name] == num_nodes - num_n:
                self.parent().vs[-num_n:][name] = values
            else:
                for n, val in zip(nodes, values):
                    self.parent().vs[n][name] = val
        self._num_values_set[name] = num_nodes


class _IgEProperty(BaseProperty):

    ''' Class for generic interactions with nodes properties (igraph) '''

    def __getitem__(self, name):
        if isinstance(name, slice):
            eprop = {}
            for k in self.keys():
                eprop[k] = np.array(self.parent().es[k])[name]
            return eprop
        elif nonstring_container(name):
            eprop = {}
            if nonstring_container(name[0]):
                eids = [self.parent().get_eid(*e) for e in name]
                for k in self.keys():
                    eprop[k] = np.array(self.parent().es[k])[eids]
            else:
                eid = self.parent().get_eid(*name)
                for k in self.keys():
                    eprop[k] = self.parent().es[k][eid]
            return eprop
        return np.array(self.parent().es[name])

    def __setitem__(self, name, value):
        if name in self:
            size = self.parent().edge_nb()
            if len(value) == size:
                self.parent().es[name] = value
            else:
                raise ValueError("A list or a np.array with one entry per "
                                 "edge in the graph is required")
        else:
            raise InvalidArgument("Attribute does not exist yet, use "
                                  "set_attribute to create it.")

    def new_attribute(self, name, value_type, values=None, val=None):
        if val is None:
            if value_type == "int":
                val = int(0)
            elif value_type == "double":
                val = 0.
            elif value_type == "string":
                val = ""
            else:
                val = None
                value_type = 'object'
        if values is None:
            values = np.repeat(val, self.parent().ecount())
        # store name and value type in the dict
        super(_IgEProperty, self).__setitem__(name, value_type)
        # store the real values in the attribute
        self[name] = values
        self._num_values_set[name] = len(values)

    def set_attribute(self, name, values, edges=None):
        '''
        Set the edge property.
        
        Parameters
        ----------
        name : str
            Name of the edge property.
        values : array
            Values that should be set.
        edges : array-like, optional (default: None)
            Edges for which the value of the property should be set. If `edges`
            is not None, it must be an array of shape `(len(values), 2)`.
        '''
        num_edges = self.parent().ecount()
        num_e = len(edges) if edges is not None else num_edges
        if num_e == num_edges:
            self[name] = values
        else:
            if num_e != len(values):
                raise ValueError("`edges` and `values` must have the same "
                                 "size; got respectively " + str(num_e) + \
                                 " and " + str(len(values)) + "entries.")
            if self._num_values_set[name] == num_edges - num_e:
                self.parent().es[-num_e:][name] = values
            else:
                for e, val in zip(edges, values):
                    eid = self.parent().get_eid(*e)
                    self.parent().es[eid][name] = val
        self._num_values_set[name] = num_edges


#-----------------------------------------------------------------------------#
# Graph
#------------------------
#

class _IGraph(BaseGraph):

    '''
    Subclass of :class:`igraph.Graph`.
    '''
    
    nattr_class = _IgNProperty
    eattr_class = _IgEProperty

    #-------------------------------------------------------------------------#
    # Constructor and instance properties
    
    def __init__(self, nodes=0, g=None, directed=True, weighted=False):
        self._nattr = _IgNProperty(self)
        self._eattr = _IgEProperty(self)
        self._weighted = weighted
        self._directed = directed
        if g is None:
            super(_IGraph,self).__init__(n=nodes, directed=True)
        else:
            nodes = g.vcount()
            edges = g.ecount()
            di_node_attr = {}
            di_edge_attr = {}
            if nodes:
                nattr = g.vs[0].attributes().keys()
            if edges:
                eattr = g.es[0].attributes().keys()
            for attr in nattr:
                di_node_attr[attr] = np.array(g.vs[:][attr])
            for attr in eattr:
                di_edge_attr[attr] = np.array(g.es[:][attr])
            super(_IGraph, self).__init__(
                n=nodes, directed=True, vertex_attrs=di_node_attr)
            lst_edges = nngt.analyze_graph["get_edges"](g)
            self.new_edges(lst_edges, attributes=di_edge_attr)

    #-------------------------------------------------------------------------#
    # Graph manipulation

    def edge_id(self, edge):
        '''
        Return the ID a given edge or a list of edges in the graph.
        Raises an error if the edge is not in the graph or if one of the
        vertices in the edge is nonexistent.

        Parameters
        ----------
        edge : 2-tuple or array of edges
            Edge descriptor (source, target).

        Returns
        -------
        index : int or array of ints
            Index of the given `edge`.
        '''
        if isinstance(edge[0], int):
            return self.get_eid(*edge)
        elif hasattr(edge[0], "__len__"):
            return self.get_eids(edge)
        else:
            raise AttributeError("`edge` must be either a 2-tuple of ints or\
an array of 2-tuples of ints.")
    
    @property
    def edges_array(self):
        ''' Edges of the graph, sorted by order of creation, as an array of
        2-tuple. '''
        return np.array([(e.source, e.target) for e in self.es])
    
    def new_node(self, n=1, ntype=1, attributes=None, value_types=None):
        '''
        Adding a node to the graph, with optional properties.
        
        Parameters
        ----------
        n : int, optional (default: 1)
            Number of nodes to add.
        ntype : int, optional (default: 1)
            Type of neuron (1 for excitatory, -1 for inhibitory)
            
        Returns
        -------
        The node or an iterator over the nodes created.
        '''
        first_node_idx = self.vcount()
        super(_IGraph, self).add_vertices(n)
        nodes = list(range(first_node_idx, first_node_idx + n))
        if attributes is not None:
            for k, v in attributes.items():
                if k not in self._nattr:
                    self._nattr.new_attribute(k, value_types[k], val=v)
                else:
                    v = v if nonstring_container(v) else [v]
                    self._nattr.set_attribute(k, v, nodes=nodes)
        self.vs[nodes[0]:nodes[-1] + 1]['type'] = ntype
        if n == 1:
            return nodes[0]
        return nodes

    def new_edge(self, source, target, attributes=None, ignore=False):
        '''
        Adding a connection to the graph, with optional properties.
        
        Parameters
        ----------
        source : :class:`int/node`
            Source node.
        target : :class:`int/node`
            Target node.
        attributes : :class:`dict`, optional (default: ``{}``)
            Dictionary containing optional edge properties. If the graph is
            weighted, defaults to ``{"weight": 1.}``, the unit weight for the
            connection (synaptic strength in NEST).
        ignore : bool, optional (default: False)
            If set to True, ignore attempts to add an existing edge, otherwise
            raises an error.
            
        Returns
        -------
        The new connection.
        '''
        if attributes is not None:
            attributes = {k: [v] for k, v in attributes.items()}
        self.new_edges(((source, target),), attributes)

    def new_edges(self, edge_list, attributes=None):
        '''
        Add a list of edges to the graph.
        
        .. warning ::
            This function currently does not check for duplicate edges!
        
        Parameters
        ----------
        edge_list : list of 2-tuples or np.array of shape (edge_nb, 2)
            List of the edges that should be added as tuples (source, target)
        attributes : dict, optional (default: ``None``)
            Dictionary of the form ``{"name": [], "values": [],
            "type": []}``, containing the attributes of the new edges.
            
        @todo: add example, check the edges for self-loops and multiple edges
        
        Returns
        -------
        Returns new edges only.
        '''
        if attributes is None:
            attributes = {}
        initial_ecount = self.ecount()
        edge_list = np.array(edge_list)
        if not self._directed:
            recip_edges = edge_list[:,::-1]
            # slow but works
            unique = ~(recip_edges[..., np.newaxis]
                      == edge_list[..., np.newaxis].T).all(1).any(1)
            edge_list = np.concatenate((edge_list, recip_edges[unique]))
            for key, val in attributes.items():
                attributes[key] = np.concatenate((val, val[unique]))
        first_eid = self.ecount()
        super(_IGraph, self).add_edges(edge_list)
        # call parent function to set the attributes
        self.attr_new_edges(edge_list, attributes=attributes)
        return edge_list

    def remove_edge(self, edge):
        raise NotImplementedError("This function has been removed because it \
            makes using edge properties too complicated")

    def remove_vertex(self, node, fast=False):
        raise NotImplementedError("This function has been removed because it \
            makes using node properties too complicated")

    def clear_all_edges(self):
        ''' Remove all connections in the graph. '''
        self.delete_edges(None)
        self._eattr.clear()
    
    #-------------------------------------------------------------------------#
    # Getters
    
    def node_nb(self):
        return self.vcount()

    def edge_nb(self):
        return self.ecount()
    
    def degree_list(self, node_list=None, deg_type="total", use_weights=False):
        deg_type = 'all' if deg_type == 'total' else deg_type
        if use_weights:
            return np.array(self.strength(node_list, mode=deg_type, weights='weight'))
        else:
            return np.array(self.degree(node_list, mode=deg_type))

    def betweenness_list(self, btype="both", use_weights=False, norm=True,
                         **kwargs):
        n = self.vcount()
        e = self.ecount()
        ncoeff_norm = (n-1)*(n-2)
        ecoeff_norm = (e-1)*(e-2)/2.
        w, nbetw, ebetw = None, None, None
        if use_weights:
            w = self.es['bweight']
        if btype in ("both", "node"):
            nbetw = np.array(self.betweenness(weights=w))
        if btype in ("both", "edge"):
            ebetw = np.array(self.edge_betweenness(weights=w))
        if btype == "node":
            return nbetw/ncoeff_norm if norm else nbetw
        elif btype == "edge":
            return ebetw/ecoeff_norm if norm else ebetw
        else:
            return ( nbetw/ncoeff_norm, ebetw/ecoeff_norm if norm
                     else nbetw, ebetw )

    def neighbours(self, node, mode="all"):
        '''
        Return the neighbours of `node`.

        Parameters
        ----------
        node : int
            Index of the node of interest.
        mode : string, optional (default: "all")
            Type of neighbours that will be returned: "all" returns all the
            neighbours regardless of directionality, "in" returns the
            in-neighbours (also called predecessors) and "out" retruns the
            out-neighbours (or successors).

        Returns
        -------
        neighbours : tuple
            The neighbours of `node`.
        '''
        if mode == "all":
            return self.neighbors(node, mode=3)
        elif mode == "in":
            return self.neighbors(node, mode=1)
        elif mode == "out":
            return self.neighbors(node, mode=2)
        else:
            raise ArgumentError('''Invalid `mode` argument {}; possible values
                                are "all", "out" or "in".'''.format(mode))

    #-------------------------------------------------------------------------
    # Prevent users from calling igraph functions

    def add_vertex(self, *args, **kwargs):
        raise RuntimeError("Intrinsic igraph functions for vertex "
                           "creation have been disabled.")

    def add_vertices(self, *args, **kwargs):
        raise RuntimeError("Intrinsic igraph functions for vertex "
                           "creation have been disabled.")

    def add_edge(self, *args, **kwargs):
        raise RuntimeError("Intrinsic igraph functions for edge "
                           "creation have been disabled.")

    def add_edges(self, *args, **kwargs):
        raise RuntimeError("Intrinsic igraph functions for edge "
                           "creation have been disabled.")

    def clear_edges(self, *args, **kwargs):
        raise RuntimeError("Intrinsic graph_tool functions for edge "
                           "deletion have been disabled.")
