# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
# -----------------------------------------------------------------------------
"""Windbarb Visual and shader definitions."""

import numpy as np

from vispy.color import ColorArray
from vispy.gloo import VertexBuffer
from vispy.visuals.shaders import Variable
from vispy.visuals.visual import Visual

_VERTEX_SHADER = """
uniform float u_antialias;
uniform float u_px_scale;
uniform float u_scale;

attribute vec3  a_position;
attribute vec2  a_wind;
attribute vec4  a_fg_color;
attribute vec4  a_bg_color;
attribute float a_edgewidth;
attribute float a_size;
attribute float a_trig;

varying vec4 v_fg_color;
varying vec4 v_bg_color;
varying vec2 v_wind;
varying float v_trig;
varying float v_edgewidth;
varying float v_antialias;

void main (void) {
    $v_size = a_size * u_px_scale * u_scale;
    v_edgewidth = a_edgewidth * float(u_px_scale);
    v_wind = a_wind.xy;
    v_trig = a_trig;
    v_antialias = u_antialias;
    v_fg_color  = a_fg_color;
    v_bg_color  = a_bg_color;
    gl_Position = $transform(vec4(a_position,1.0));
    float edgewidth = max(v_edgewidth, 1.0);
    gl_PointSize = ($v_size) + 4.*(edgewidth + 1.5*v_antialias);
}
"""

_FRAGMENT_SHADER = """
#include "math/constants.glsl"
#include "math/signed-segment-distance.glsl"
#include "antialias/antialias.glsl"
varying vec4 v_fg_color;
varying vec4 v_bg_color;
varying vec2 v_wind;
varying float v_trig;
varying float v_edgewidth;
varying float v_antialias;

// SDF-Triangle by @rougier
// https://github.com/rougier/python-opengl/blob/master/code/chapter-06/SDF-triangle.py
float sdf_triangle(vec2 p, vec2 p0, vec2 p1, vec2 p2)
{
    vec2 e0 = p1 - p0;
    vec2 e1 = p2 - p1;
    vec2 e2 = p0 - p2;
    vec2 v0 = p - p0;
    vec2 v1 = p - p1;
    vec2 v2 = p - p2;
    vec2 pq0 = v0 - e0*clamp( dot(v0,e0)/dot(e0,e0), 0.0, 1.0 );
    vec2 pq1 = v1 - e1*clamp( dot(v1,e1)/dot(e1,e1), 0.0, 1.0 );
    vec2 pq2 = v2 - e2*clamp( dot(v2,e2)/dot(e2,e2), 0.0, 1.0 );
    float s = sign( e0.x*e2.y - e0.y*e2.x );
    vec2 d = min( min( vec2( dot( pq0, pq0 ), s*(v0.x*e0.y-v0.y*e0.x) ),
                     vec2( dot( pq1, pq1 ), s*(v1.x*e1.y-v1.y*e1.x) )),
                     vec2( dot( pq2, pq2 ), s*(v2.x*e2.y-v2.y*e2.x) ));
    return -sqrt(d.x)*sign(d.y);
}

void main()
{
    // Discard plotting marker body and edge if zero-size
    if ($v_size <= 0.)
        discard;

    float edgewidth = max(v_edgewidth, 1.0);
    float linewidth = max(v_edgewidth, 1.0);
    float edgealphafactor = min(v_edgewidth, 1.0);

    float size = $v_size + 4.*(edgewidth + 1.5*v_antialias);
    // factor 6 for acute edge angles that need room as for star marker
    
    vec2 wind = v_wind;
    
    if (v_trig > 0.)
    {
        float u = wind.x * cos(radians(wind.y));
        float v = wind.x * sin(radians(wind.y));
        wind = vec2(u, v);
    }
    
    // knots to m/s
    wind *= 2.;
    
    // normalized distance
    float dx = 0.5;
    // normalized center point
    vec2 O = vec2(dx);
    // normalized x-component
    vec2 X = normalize(wind) * dx / M_SQRT2 / 1.1 * vec2(1, -1);
    // normalized y-component
    // here the barb can be mirrored for southern earth * (vec2(1., -1.)
    //vec2 Y = X.yx * vec2(1., -1.); // southern hemisphere
    vec2 Y = X.yx * vec2(-1., 1.); // northern hemisphere
    // PointCoordinate
    vec2 P = gl_PointCoord;

    // calculate barb items
    float speed = length(wind);
    int flag = int(floor(speed / 50.));
    speed -= float (50 * flag);
    int longbarb = int(floor(speed / 10.));
    speed -= float (longbarb * 10);
    int shortbarb = int(floor(speed / 5.));
    int calm = shortbarb + longbarb + flag;

    // starting distance
    float r;
    // calm, plot circles
    if (calm == 0)
    {
        r = abs(length(O-P)- dx * 0.2);
        r = min(r, abs(length(O-P)- dx * 0.1));
    }
    else
    {
        // plot shaft
        r = segment_distance(P, O, O-X);
        float pos = 1.;

        // plot flag(s)
        while(flag >= 1)
        {
            r = min(r, sdf_triangle(P, O-X*pos, O-X*pos-X*.4-Y*.4, O-X*pos-X*.4));
            flag -= 1;
            pos -= 0.15;
        }
        // plot longbarb(s)
        while(longbarb >= 1)
        {
            r = min(r, segment_distance(P, O-X*pos, O-X*pos-X*.4-Y*.4));
            longbarb -= 1;
            pos -= 0.15;
        }
        // plot shortbarb
        while(shortbarb >= 1)
        {
            if (pos == 1.0)
                pos -= 0.15;
            r = min(r, segment_distance(P, O-X*pos, O-X*pos-X*.2-Y*.2));
            shortbarb -= 1;
            pos -= 0.15;
        }
    }

    // apply correction for size
    r *= size;

    vec4 edgecolor = vec4(v_fg_color.rgb, edgealphafactor*v_fg_color.a);

    if (r > 0.5 * v_edgewidth + v_antialias)
    {
        // out of the marker (beyond the outer edge of the edge
        // including transition zone due to antialiasing)
        discard;
    }

    gl_FragColor = filled(r, edgewidth, v_antialias, edgecolor);
}
"""


class WindbarbVisual(Visual):
    """Visual displaying windbarbs."""

    _shaders = {
        'vertex': _VERTEX_SHADER,
        'fragment': _FRAGMENT_SHADER,
    }

    def __init__(self, **kwargs):
        self._vbo = VertexBuffer()
        self._v_size_var = Variable('varying float v_size')
        self._marker_fun = None
        self._data = None
        Visual.__init__(self, vcode=self._shaders['vertex'], fcode=self._shaders['fragment'])
        self.shared_program.vert['v_size'] = self._v_size_var
        self.shared_program.frag['v_size'] = self._v_size_var
        self.set_gl_state(depth_test=True, blend=True,
                          blend_func=('src_alpha', 'one_minus_src_alpha'))
        self._draw_mode = 'points'
        if len(kwargs) > 0:
            self.set_data(**kwargs)
        self.freeze()

    def set_data(self, pos=None, wind=None, trig=True, size=50.,
                 antialias=1., edge_width=1., edge_color='black',
                 face_color='white'):
        """Set the data used to display this visual.

        Parameters
        ----------
        pos : array
            The array of locations to display each windbarb.
        wind : array
            The array of wind vector components to display each windbarb.
            in m/s. For knots divide by two.
        trig : bool
            True - wind contains (mag, ang)
            False - wind contains (u, v)
            defaults to True
        size : float or array
            The windbarb size in px.
        antialias : float
            The antialiased area (in pixels).
        edge_width : float | None
            The width of the windbarb outline in pixels.
        edge_color : Color | ColorArray
            The color used to draw each symbol outline.
        face_color : Color | ColorArray
            The color used to draw each symbol interior.
        """
        assert (isinstance(pos, np.ndarray) and
                pos.ndim == 2 and pos.shape[1] in (2, 3))
        assert (isinstance(wind, np.ndarray) and
                pos.ndim == 2 and pos.shape[1] == 2)
        if edge_width < 0:
            raise ValueError('edge_width cannot be negative')

        # since the windbarb starts in the fragment center,
        # we need to multiply by 2 for correct length
        size *= 2

        edge_color = ColorArray(edge_color).rgba
        if len(edge_color) == 1:
            edge_color = edge_color[0]

        face_color = ColorArray(face_color).rgba
        if len(face_color) == 1:
            face_color = face_color[0]

        n = len(pos)
        data = np.zeros(n, dtype=[('a_position', np.float32, 3),
                                  ('a_wind', np.float32, 2),
                                  ('a_trig', np.float32, 0),
                                  ('a_fg_color', np.float32, 4),
                                  ('a_bg_color', np.float32, 4),
                                  ('a_size', np.float32),
                                  ('a_edgewidth', np.float32)])
        data['a_fg_color'] = edge_color
        data['a_bg_color'] = face_color
        data['a_edgewidth'] = edge_width
        data['a_position'][:, :pos.shape[1]] = pos
        data['a_wind'][:, :wind.shape[1]] = wind
        if trig:
            data['a_trig'] = 1.
        else:
            data['a_trig'] = 0.
        data['a_size'] = size
        self.shared_program['u_antialias'] = antialias
        self._data = data
        self._vbo.set_data(data)
        self.shared_program.bind(self._vbo)
        self.update()

    def _prepare_transforms(self, view):
        xform = view.transforms.get_transform()
        view.view_program.vert['transform'] = xform

    def _prepare_draw(self, view):
        view.view_program['u_px_scale'] = view.transforms.pixel_scale
        view.view_program['u_scale'] = 1

    def _compute_bounds(self, axis, view):
        pos = self._data['a_position']
        if pos is None:
            return None
        if pos.shape[1] > axis:
            return (pos[:, axis].min(), pos[:, axis].max())
        else:
            return (0, 0)