#!/usr/bin/env python
"""
This script generates a force and velocity vector diagram for a cross-flow
turbine.
"""

from __future__ import division, print_function
import numpy as np
import matplotlib
import matplotlib.pyplot as plt


def mag(v):
    """
    Return magnitude of 2-D vector (input as a tuple, list, or NumPy array).
    """
    return np.sqrt(v[0]**2 + v[1]**2)
    

def rotate(v, rad):
    """Rotate a 2-D vector by rad radians."""
    dc, ds = np.cos(rad), np.sin(rad)
    x, y = v[0], v[1]
    x, y = dc*x - ds*y, ds*x + dc*y
    return np.array((x, y))


def gen_naca_points(naca="0020", c=100, npoints=100, tuples=True):
    """Generate points for a NACA foil."""
    x = np.linspace(0, 1, npoints)*c
    t = float(naca[2:])/100.0
    y = 5.0*t*c*(0.2969*np.sqrt(x/c) - 0.1260*(x/c) - 0.3516*(x/c)**2 \
            + 0.2843*(x/c)**3 - 0.1015*(x/c)**4)
    y = np.append(y, -y[::-1])
    x = np.append(x, x[::-1])
    if tuples:
        return np.array([(x0, y0) for x0, y0 in zip(x, y)])
    else:
        return x, y


def test_gen_naca_points():
    points = gen_naca_points()
    x = []
    y = []
    for p in points:
        x.append(p[0])
        y.append(p[1])
    fig, ax = plt.subplots()
    ax.plot(x, y, "o")
    ax.set_aspect(1)
    plt.show()


def plot_radius(ax, theta_deg=0):
    """Plot radius at given azimuthal angle."""
    r = 0.495
    theta_rad = np.deg2rad(theta_deg)
    x2, y2 = r*np.cos(theta_rad), r*np.sin(theta_rad)
    ax.plot((0, x2), (0, y2), "k", linewidth=2)


def make_naca_path(c=0.3, theta_deg=0.0):
    verts = gen_naca_points(c=c)
    verts = np.array([rotate(v, -np.pi/2) for v in verts])
    verts += (0.5, c/4)
    theta_rad = np.deg2rad(theta_deg)
    verts = np.array([rotate(v, theta_rad) for v in verts])
    p = matplotlib.path.Path(verts, closed=True)
    return p
        
    
def plot_foil(ax, c=0.3, theta_deg=0.0):
    """Plot the foil shape using a matplotlib patch."""
    p = matplotlib.patches.PathPatch(make_naca_path(c, theta_deg), 
                                     facecolor="gray", linewidth=1)
    ax.add_patch(p)


def plot_velocities(ax, theta_deg=0.0, tsr=2.0):
    """Plot blade velocity, free stream velocity, and relative velocity."""
    r = 0.5
    u_infty = 0.25
    theta_rad = np.deg2rad(theta_deg)
    blade_xy = r*np.cos(theta_rad), r*np.sin(theta_rad)
    head_width = 0.05
    head_length = 0.12
    linewidth = 1.5
    
    # Make blade velocity vector
    x1, y1 = rotate((0.5, tsr*u_infty), np.deg2rad(theta_deg))
    print(blade_xy)
    dx, dy = np.array(blade_xy) - np.array((x1, y1))
    print(dx, dy)
    ax.arrow(x1, y1, dx, dy, head_width=head_width, head_length=head_length, 
             length_includes_head=True, color="black", linewidth=linewidth)
    ax.text(x1 + 0.01, y1 + 0.05*np.sign(y1), r"$\omega r$")
    
    # Make free stream velocity vector
    y1 += u_infty
    ax.arrow(x1, y1, 0, -u_infty, head_width=head_width, 
             head_length=head_length, length_includes_head=True, color="blue",
             linewidth=linewidth)


def plot_diagram(theta_deg=0.0, tsr=2.0, save=False):
    """Plot full vector diagram."""
    fig, ax = plt.subplots(figsize=(6, 6))
    
    plot_foil(ax, c=0.3, theta_deg=theta_deg)
    plot_radius(ax, theta_deg)
    plot_velocities(ax, theta_deg, tsr)

    # Figure formatting    
    ax.set_xlim((-1, 1))
    ax.set_ylim((-1, 1))
    ax.set_aspect(1)
    ax.axis("off")
    
    if save:
        fig.savefig("cft-vectors.pdf")


if __name__ == "__main__":
    plt.rcParams["font.size"] = 18
    plot_diagram(90)

@@ -1491,17 +1491,20 @@
 , y2), %22
-k
+gray
 %22, linew
@@ -2076,16 +2076,71 @@
 ewidth=1
+,%0A                                     edgecolor=%22gray%22
 )%0A    ax
@@ -2312,10 +2312,9 @@
 = 0.
-25
+3
 %0A   
@@ -2574,28 +2574,8 @@
 g))%0A
-    print(blade_xy)%0A
     
@@ -2627,26 +2627,8 @@
 1))%0A
-    print(dx, dy)%0A
     
@@ -3071,16 +3071,273 @@
 ewidth)%0A
+             %0A    # Make relative velocity vector%0A    dx, dy = np.array(blade_xy) - np.array((x1, y1))%0A    ax.arrow(x1, y1, dx, dy, head_width=head_width, head_length=head_length, %0A             length_includes_head=True, color=%22green%22, linewidth=linewidth)%0A
 %0A%0Adef pl
@@ -3842,12 +3842,13 @@
 diagram(
-9
+30
 0)%0A

Add relative velocity vector