from .rxdException import RxDException

class NodeList(list):
    def __init__(self, items):
        """Constructs a NodeList from items, a python iterable containing Node objects."""
        list.__init__(self, items)
    def __call__(self, restriction):
        """returns a sub-NodeList consisting of nodes satisfying restriction"""
        return NodeList([i for i in self if i.satisfies(restriction)])
    
    @property
    def value(self):
        # TODO: change this when not everything is a concentration
        return self.concentration
    @value.setter
    def value(self, v):
        # TODO: change this when not everything is a concentration
        self.concentration = v
    
    @property
    def concentration(self):
        """Returns the concentration of the Node objects in the NodeList as an iterable."""
        return [node.concentration for node in self]
    @concentration.setter
    def concentration(self, value):
        """Sets the concentration of the Node objects to either a constant or values based on an iterable."""
        if hasattr(value, '__len__'):
            if len(value) == len(self):
                for node, val in zip(self, value): node.concentration = val
                return
            else:
                raise RxDException('concentration must either be a scalar or an iterable of the same length as the NodeList')
        for node in self: node.concentration = value
    
    @property
    def diff(self):
        """Returns the diffusion constant of the Node objects in the NodeList as an iterable."""
        return [node.diff for node in self]
    @diff.setter
    def diff(self, value):
        """Sets the diffusion constant of the Node objects to either a constant or values based on an iterable."""
        if hasattr(value, '__len__'):
            if len(value) == len(self):
                for node, val in zip(self, value): node.diff = val
                return
            else:
                raise RxDException('diff must either be a scalar or an iterable of the same length as the NodeList')
        for node in self: node.diff = value
        
    def value_to_grid(self):
        """Returns a regular grid with the values of the 3d nodes in the list.
        
        The grid is a copy only.
        
        Grid points not belonging to the object are assigned a value of NaN.
        
        Nodes that are not 3d will be ignored. If there are no 3d nodes, returns
        a 0x0x0 numpy array.
        
        Warning: Currently only supports nodelists over 1 region.
        """
        import numpy
        from .node import Node3D
        
        # identify regions involved
        regions = set()
        for node in self:
            if isinstance(node, Node3D):
                regions.add(node.region)
        
        # default result that falls through if no regions
        result = numpy.zeros((0, 0, 0))

        if len(regions) > 1:
            # TODO: the reason for this restriction is the need to make sure
            #       the mesh lines up
            raise RxDException('value_to_grid currently only supports 1 region')
        
        for r in regions:
            # TODO: if allowing multiple regions, need to change this
            result = numpy.empty(r._mesh.shape)
            result.fill(numpy.nan)
            
            for node in self:
                if isinstance(node, Node3D):
                    # TODO: if allowing multiple regions, adjust i, j, k as needed
                    result[node._i, node._j, node._k] = node.value
        
        return result
            
        

    @property
    def volume(self):
        """An iterable of the volumes of the Node objects in the NodeList.

        Read only."""
        return [node.volume for node in self]

    @property
    def surface_area(self):
        """An iterable of the surface areas of the Node objects in the NodeList.

        Read only."""
        return [node.surface_area for node in self]

    @property
    def region(self):
        """An iterable of the Region objects corresponding to the Node objects in the NodeList.

        Read only."""
        return [node.region for node in self]
    
    @property
    def species(self):
        """An iterable of the Species objects corresponding to the Node objects in the NodeList.

        Read only."""
        return [node.species for node in self]
    
    @property
    def x(self):
        """An iterable of the normalized positions of the Node objects in the NodeList.

        Read only."""
        return [node.x for node in self]
