import itertools
import numpy as np
from ... import geometry
from ... import mesh
from .internal import GLPrimitive, GLShapeDecorator
from ... import draw
from ..internal import ShapeAttribute
from vispy import gloo

@GLShapeDecorator
class Polygons(draw.Polygons, GLPrimitive):
    __doc__ = draw.Polygons.__doc__

    shaders = {}

    shaders['vertex'] = """
       uniform mat4 camera;
       uniform vec4 rotation;
       uniform vec3 translation;
       uniform float outline;

       attribute vec4 color;
       attribute vec2 position;
       attribute vec2 image;
       attribute vec2 outline_image;
       attribute vec4 orientation;

       varying vec4 v_color;

       vec2 rotate(vec2 point, vec4 quat)
       {
           vec3 point3d = vec3(point.xy, 0.0);
           vec3 result = (quat.x*quat.x - dot(quat.yzw, quat.yzw))*point3d;
           result += 2.0*quat.x*cross(quat.yzw, point3d);
           result += 2.0*dot(quat.yzw, point3d)*quat.yzw;
           return result.xy;
       }

       void main()
       {
           vec2 currentImage = mix(image, outline_image, max(0.0, outline));
           vec4 currentColor = color;
           if(outline < 0.0)
               currentColor = vec4(0.0, 0.0, 0.0, color.a);

           vec2 vertexPos = position + rotate(currentImage, orientation);
           vertexPos = rotate(vertexPos, rotation) + translation.xy;
           vec4 screenPosition = camera * vec4(vertexPos, translation.z, 1.0);

           // transform to screen coordinates
           gl_Position = screenPosition;
           v_color = currentColor;
       }
       """

    shaders['fragment'] = """
       varying vec4 v_color;

       void main()
       {
           gl_FragColor = v_color;
       }
       """

    _vertex_attribute_names = ['position', 'orientation', 'color', 'image', 'outline_image']

    _GL_UNIFORMS = list(itertools.starmap(ShapeAttribute, [
        ('camera', np.float32, np.eye(4), 2,
         'Internal: 4x4 Camera matrix for world projection'),
        ('rotation', np.float32, (1, 0, 0, 0), 1,
         'Internal: Rotation to be applied to each scene as a quaternion'),
        ('translation', np.float32, (0, 0, 0), 1,
         'Internal: Translation to be applied to the scene'),
        ('outline', np.float32, 0, 0,
         'Outline width for shapes')
        ]))

    def __init__(self, *args, **kwargs):
        GLPrimitive.__init__(self)
        draw.Polygons.__init__(self, *args, **kwargs)

    def update_arrays(self):
        if 'vertices' in self._dirty_attributes:
            vertices = self.vertices
            if len(vertices) < 3:
                thetas = np.linspace(0, 2*np.pi, 3, endpoint=False)
                vertices = np.array([np.cos(thetas), np.sin(thetas)], dtype=np.float32).T

            polygon = geometry.Polygon(vertices)
            self._gl_attributes['triangulation'] = geometry.Outline(polygon, 1.0)
            self._gl_attributes['indices'] = self._gl_attributes['triangulation'].outer.triangleIndices

        try:
            for name in self._dirty_attributes:
                if name == 'vertices':
                    raise ValueError
                else:
                    self._gl_vertex_arrays[name][:] = self._attributes[name]
                    self._dirty_vertex_attribs.add(name)
        except (ValueError, KeyError):
            vertices = self._gl_attributes['triangulation'].outer.vertices
            outline_vertices = self._gl_attributes['triangulation'].inner.vertices

            vertex_arrays = mesh.unfoldProperties(
                [self.positions, self.orientations, self.colors],
                [vertices, outline_vertices])

            unfolded_shape = vertex_arrays[0].shape[:-1]
            indices = (np.arange(unfolded_shape[0])[:, np.newaxis, np.newaxis]*unfolded_shape[1] +
                       self._gl_attributes['indices'])
            indices = indices.reshape((-1, 3))

            self._finalize_array_updates(indices, vertex_arrays)

        self._dirty_attributes.clear()

    def render_generic(self, *args, **kwargs):
        try:
            outline = self.outline.copy()
            translation = self.translation.copy()
            super(Polygons, self).render_generic(*args, **kwargs)
            if outline > 0:
                self.outline = -1
                self.translation = translation - (0, 0, 1e-3)
                super(Polygons, self).render_generic(*args, **kwargs)
        finally:
            self.outline = outline
            self.translation = translation
