import numpy as np
from colorsys import hsv_to_rgb
from collections import deque

from .util      import is_sequence, is_shape, tracked_array, Cache, DataStore
from .constants import log

COLORS = {'red'    : [205,59,34],
          'purple' : [150,111,214],
          'blue'   : [119,158,203],
          'brown'  : [160,85,45]}
COLOR_DTYPE = np.dtype(np.uint8)
DEFAULT_COLOR = COLORS['purple']

class VisualAttributes(object):
    '''
    Hold the visual attributes (usually colors) for a mesh. 
    '''
    def __init__(self, mesh=None, **kwargs):
        self.mesh = mesh

        self._validate = True

        self._data = DataStore()
        # cache computed values which are cleared when
        # self.md5() changes, forcing a recompute
        self._cache = Cache(id_function = self._data.md5)

        colors = _kwargs_to_color(mesh, **kwargs)
        self.vertex_colors, self.face_colors = colors



    @property
    def _set(self):
        result = {'face'   : is_shape(self._data['face_colors'], (-1, (3,4))),
                  'vertex' : is_shape(self._data['vertex_colors'], (-1,(3,4)))}
        return result

    @property
    def defined(self):
        defined = np.any(self._set.values()) and self.mesh is not None
        return defined

    @property
    def transparency(self):
        cached = self._cache.get('transparency')
        if cached is not None: 
            return cached
        transparency = False
        color_max = (2**(COLOR_DTYPE.itemsize*8)) - 1
        if self._set['face']:
            transparency = (is_shape(self._data['face_colors'], (-1,4)) and 
                            np.any(self._data['face_colors'][:,3] < color_max))
        elif self._set['vertex']:
            transparency = (self._vertex_colors.ndim == 2 and
                            self._vertex_colors.shape[1] == 4 and
                            np.any(self._vertex_colors[:,3] < color_max))

        return self._cache.set(key   = 'transparency',
                               value = bool(transparency))

    def md5(self):
        return self._data.md5()

    @property
    def face_colors(self):
        cached = self._data['face_colors']
        if not (is_sequence(cached) and
                len(cached) == len(self.mesh.faces)):
            log.debug('Returning default colors for faces')
            return np.tile(DEFAULT_COLOR, (len(self.mesh.faces), 1))
        return cached

    @face_colors.setter
    def face_colors(self, values):
        values = np.asanyarray(values)
        if values.shape in ((3,), (4,)):
            # case where a single RGB/RGBa color has been passed to the setter
            # we apply this color to all faces 
            values = np.tile(values, (len(self.mesh.faces), 1))
        self._data['face_colors'] = values

    @property
    def vertex_colors(self):
        stored = self._data['vertex_colors']
        if not (is_sequence(stored) and
                len(stored) == len(self.mesh.vertices)):
            log.debug('Vertex colors being generated from face colors')
            return face_to_vertex_color(self.mesh, self.face_colors)
        return stored

    @vertex_colors.setter
    def vertex_colors(self, values):
        self._data['vertex_colors'] = values

    def update_faces(self, mask):
        stored = self._data['face_colors']
        if not is_shape(stored, (-1,(3,4))):
            return 
        try: 
            self._data['face_colors'] = stored[mask]
        except: 
            log.warning('Face colors not updated', exc_info=True) 
            
    def update_vertices(self, mask):
        stored = self._data['vertex_colors']
        if not is_shape(stored, (-1, (3,4))):
            return
        try:    
            self._data['vertex_colors'] = stored[mask]
        except: 
            log.debug('Vertex colors not updated', exc_info=True)

    def subsets(self, faces_sequence):
        result = deque()
        face   = self.face_colors
        for f in faces_sequence:
            result.append(VisualAttributes())
            result[-1].face_colors = face[list(f)]
        return np.array(result)

    def union(self, others):
        return visuals_union(np.append(self, others))

def _kwargs_to_color(mesh, **kwargs):
    '''
    Given a set of keyword arguments, see if any reference color
    in their name, and match the dimensions of the mesh.
    '''

    def pick_option(vf):
        if any(i is None for i in vf):
            return vf
        result = [None, None]
        signal = [i.ptp(axis=0).sum() for i in vf]
        signal_max = np.argmax(signal)
        result[signal_max] = vf[signal_max]
        return result
    def pick_color(sequence):
        if len(sequence) == 0:
            return None
        elif len(sequence) == 1:
            return sequence[0]
        else:
            signal = [i.ptp(axis=0).sum() for i in sequence]
            signal_max = np.argmax(signal)
            return sequence[signal_max]

    if mesh is None:
        result = [None, None]
        if 'face_colors' in kwargs:
            result[1] = np.asanyarray(kwargs['face_colors'])
        if 'vertex_colors' in kwargs:
            result[0] = np.asanyarray(kwargs['vertex_colors'])
        return result
        
    vertex = deque()
    face   = deque()
    
    for key in kwargs.keys():
        if not ('color' in key): 
            continue
        value = np.asanyarray(kwargs[key])
        if len(value) == len(mesh.vertices):
            vertex.append(value)
        elif len(value) == len(mesh.faces):
            face.append(value)
    return pick_option([pick_color(i) for i in [vertex, face]])

def visuals_union(visuals, *args):
    visuals = np.append(visuals, args)
    color = {'face_colors'   : None,
             'vertex_colors' : None}

    vertex_ok = True
    vertex = [None] * len(visuals)

    face_ok = True
    face = [None] * len(visuals)

    for i, v in enumerate(visuals):
        if v.mesh is None:
            if face_ok and v._set['face']: 
                face[i] = rgba(v._data['face_colors'])
            else:
                face_ok = False
            if vertex_ok and v._set['vertex']:
                vertex[i] = rgba(v._data['vertex_colors'])
            else:
                vertex_ok = False
        else:
            vertex[i] = rgba(v.vertex_colors)
            face[i] = rgba(v.face_colors)
    if face_ok:
        color['face_colors'] = np.vstack(face)
    if vertex_ok:
        color['vertex_colors'] = np.vstack(vertex)

    return VisualAttributes(**color)

def rgba(colors, dtype=COLOR_DTYPE):
    colors = np.asanyarray(colors)
    if is_shape(colors, (-1,3)):
        opaque = (2**(np.dtype(dtype).itemsize * 8)) - 1
        colors = np.column_stack((colors,
                                  opaque * np.ones(len(colors))))
    return colors.astype(dtype)

def random_color(dtype=COLOR_DTYPE):
    '''
    Return a random RGB color using datatype specified.
    '''
    hue  = np.random.random() + .61803
    hue %= 1.0
    color = np.array(hsv_to_rgb(hue, .99, .99))
    if np.dtype(dtype).kind in 'iu':
        max_value = (2**(np.dtype(dtype).itemsize * 8)) - 1
        color    *= max_value
    color     = color.astype(dtype)
    return color

def face_to_vertex_color(mesh, face_colors, dtype=COLOR_DTYPE):
    '''
    Convert a set of face colors into a set of vertex colors.
    '''
    
    color_dim = np.shape(face_colors)[1]

    vertex_colors = np.zeros((len(mesh.vertices),3,color_dim))
    population    = np.zeros((len(mesh.vertices),3), dtype=np.bool)

    vertex_colors[[mesh.faces[:,0],0]] = face_colors
    vertex_colors[[mesh.faces[:,1],1]] = face_colors
    vertex_colors[[mesh.faces[:,2],2]] = face_colors

    population[[mesh.faces[:,0], 0]] = True
    population[[mesh.faces[:,1], 1]] = True
    population[[mesh.faces[:,2], 2]] = True

    # clip the population sum to 1, to avoid a division error in edge cases
    populated     = np.clip(population.sum(axis=1), 1, 3)
    vertex_colors = vertex_colors.sum(axis=1) / populated.reshape((-1,1))
    return vertex_colors.astype(dtype)
