import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import my_config as config

import logging
from src.getDf import readCsvFiles, addFundPerf, readStatics, getFeatureDf, dropOff, dropOffTrade
from src.calStatics import calStatics
from src.getQuantDf import getQuantDf
from src.calQuantCode import calQuantDf

import numpy as np
import pandas as pd

import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__file__)

def debug(df):
    print df.columns.values

def mergeDailyWeekly(d_quant, w_quant):
    daily_size = d_quant.shape[0]
    quant = pd.merge(d_quant, w_quant, on='DATE', how='outer')

    #Find the first weekly is Not NaN
    drop_index = 0
    for drop_index, row in quant.iterrows():
        if not pd.isnull(row['SPY_weekly']):
            break

    quant.drop(quant.index[0:drop_index], inplace=True)
    quant = quant.reset_index(drop=True)

    print quant.columns.names

    temp_dict = {}
    for index, row in quant.iterrows():
        print index
        print row
        for name in w_quant.columns.values:
            if pd.isnull(row[name]):
                row[name] = temp_dict[name]
            else:
                temp_dict[name] = row[name]
    quant.drop(quant.index[daily_size:], inplace=True)
    quant = quant.reset_index(drop=True)

    return quant

def trim(df):
    index = 0
    start_index = 0
    end_index = df.shape[0]
    for index, row in df.iterrows():
        date = row['date']
        if date >= config.START_DATE:
            start_index = index
            break

    for index, row in df.iterrows():
        date = row['date']
        if date >= config.END_DATE:
            end_index = index
            break

    df = df[start_index:end_index]
    df = df.reset_index(drop=True)
    return df

'''
Read price df, read statics data, for each record, calculate it's quantitative code
'''
if __name__ == '__main__':

    usage = "Usage: GetQuantCode, this program will calculate the quantitiave codes of trading data"
    print usage

    d, w, m = getFeatureDf()
    d_stat, w_stat = readStatics()

    d = dropOff(d, 'daily')
    print "Usage: 8 daily Features: KD, KD_SLOPE, ROC, MFI_SLOPE, HIST_MOM, MFI, KD_RANK, MFI_RANK"

    if sys.argv[1] == 'run':
        d = trim(d)
        w = trim(w)
    debug(d)
    d_quant = calQuantDf(d, d_stat, 'daily')
    debug(d_quant)

    w = dropOff(w, 'weekly')
    debug(w)

    print "Usage: 4 weekly features: KD, MFI, KD_RANK, MFI_RANK"
    w_quant = calQuantDf(w, w_stat, 'weekly')

    debug(w_quant)
    if d.shape[0] != d_quant.shape[0] or w.shape[0] != w_quant.shape[0]:
        logger.error("Wrong quant size")
        exit(1)

    quant_df = mergeDailyWeekly(d_quant, w_quant)

    d = dropOffTrade(d, quant_df.iloc[0]['DATE'])
    print "Save trade pickle to disk"
    d.to_pickle(config.TRADE_DF_PICKLE)

    print "Save trade quant pickle to disk"
    quant_df.to_pickle(config.QUANT_DF_PICKLE)

    usage = "Usage: GetQuantCode, save trade_quant to pickle, save quant_df to pickle"
    print usage


# -*- coding: utf-8 -*-
"""

Created on Wed Jan 02 13:43:44 2013

Author: Josef Perktold
"""

from __future__ import print_function
import numpy as np
import numpy.testing as npt
import statsmodels.nonparametric.api as nparam

if __name__ == '__main__':

    np.random.seed(500)
    nobs = [250, 1000][0]
    sig_fac = 1
    x = np.random.uniform(-2, 2, size=nobs)
    x.sort()
    y_true = np.sin(x*5)/x + 2*x
    y = y_true + sig_fac * (np.sqrt(np.abs(3+x))) * np.random.normal(size=nobs)

    model = nparam.KernelReg(endog=[y],
                             exog=[x], reg_type='lc',
                             var_type='c', bw='cv_ls',
                             defaults=nparam.EstimatorSettings(efficient=True))

    sm_bw = model.bw

    sm_mean, sm_mfx = model.fit()

    model1 = nparam.KernelReg(endog=[y],
                             exog=[x], reg_type='lc',
                             var_type='c', bw='cv_ls')
    mean1, mfx1 = model1.fit()

    model2 = nparam.KernelReg(endog=[y],
                             exog=[x], reg_type='ll',
                             var_type='c', bw='cv_ls')

    mean2, mfx2 = model2.fit()

    print(model.bw)
    print(model1.bw)
    print(model2.bw)

    import matplotlib.pyplot as plt
    fig = plt.figure()
    ax = fig.add_subplot(1,1,1)
    ax.plot(x, y, 'o', alpha=0.5)
    ax.plot(x, y_true, lw=2, label='DGP mean')
    ax.plot(x, sm_mean, lw=2, label='kernel mean')
    ax.plot(x, mean2, lw=2, label='kernel mean')
    ax.legend()

    plt.show()


#!/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
import pandas as pd
from scipy.interpolate import interp1d
import seaborn as sns
from pxl.styleplot import set_sns
import os

# Define some colors (some from the Seaborn deep palette)
blue = sns.color_palette()[0]
green = sns.color_palette()[1]
dark_gray = (0.3, 0.3, 0.3)
red = sns.color_palette()[2]
purple = sns.color_palette()[3]
tan = sns.color_palette()[4]
light_blue = sns.color_palette()[5]

def load_foildata():
    """Loads NACA 0020 airfoil data at Re = 2.1 x 10^5."""
    Re = 2.1e5
    foil = "0020"
    fname = "NACA {}_T1_Re{:.3f}_M0.00_N9.0.dat".format(foil, Re/1e6)
    fpath = "data/{}".format(fname)
    alpha, cl, cd = np.loadtxt(fpath, skiprows=14, unpack=True)
    if alpha[0] != 0.0:
        alpha = np.append([0.0], alpha[:-1])
        cl = np.append([1e-12], cl[:-1])
        cd = np.append(cd[0], cd[:-1])
    # Mirror data about 0 degrees AoA since it's a symmetrical foil
    alpha = np.append(-np.flipud(alpha), alpha)
    cl = np.append(-np.flipud(cl), cl)
    cd = np.append(np.flipud(cd), cd)
    df = pd.DataFrame()
    df["alpha_deg"] = alpha
    df["cl"] = cl
    df["cd"] = cd
    return df

def lookup_foildata(alpha_deg):
    """Lookup foil characteristics at given angle of attack."""
    alpha_deg = np.asarray(alpha_deg)
    df = load_foildata()
    df["alpha_rad"] = np.deg2rad(df.alpha_deg)
    f_cl = interp1d(df.alpha_deg, df.cl, bounds_error=False)
    f_cd = interp1d(df.alpha_deg, df.cd, bounds_error=False)
    f_ct = interp1d(df.alpha_deg, df.cl*np.sin(df.alpha_rad) \
         - df.cd*np.cos(df.alpha_rad), bounds_error=False)
    cl, cd, ct = f_cl(alpha_deg), f_cd(alpha_deg), f_ct(alpha_deg)
    return {"cl": cl, "cd": cd, "ct": ct}

def calc_cft_ctorque(tsr=2.0, chord=0.14, R=0.5):
    """Calculate the geometric torque coefficient for a CFT."""
    U_infty = 1.0
    omega = tsr*U_infty/R
    theta_blade_deg = np.arange(0, 721)
    theta_blade_rad = np.deg2rad(theta_blade_deg)
    blade_vel_mag = omega*R
    blade_vel_x = blade_vel_mag*np.cos(theta_blade_rad)
    blade_vel_y = blade_vel_mag*np.sin(theta_blade_rad)
    u = U_infty # No induction
    rel_vel_mag = np.sqrt((blade_vel_x + u)**2 + blade_vel_y**2)
    rel_vel_x = u + blade_vel_x
    rel_vel_y = blade_vel_y
    relvel_dot_bladevel = (blade_vel_x*rel_vel_x + blade_vel_y*rel_vel_y)
    alpha_rad = np.arccos(relvel_dot_bladevel/(rel_vel_mag*blade_vel_mag))
    alpha_rad[theta_blade_deg > 180] *= -1
    alpha_deg = np.rad2deg(alpha_rad)
    foil_coeffs = lookup_foildata(alpha_deg)
    ctorque = foil_coeffs["ct"]*chord/(2*R)*rel_vel_mag**2/U_infty**2
    cdx = -foil_coeffs["cd"]*np.sin(np.pi/2 - alpha_rad + theta_blade_rad)
    clx = foil_coeffs["cl"]*np.cos(np.pi/2 - alpha_rad - theta_blade_rad)
    df = pd.DataFrame()
    df["theta"] = theta_blade_deg
    df["alpha_deg"] = alpha_deg
    df["rel_vel_mag"] = rel_vel_mag
    df["ctorque"] = ctorque
    df["cdrag"] = clx + cdx
    return df

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), "gray", linewidth=2)

def plot_center(ax, length=0.07, linewidth=1.2):
    """Plot centermark at origin."""
    ax.plot((0, 0), (-length/2, length/2), lw=linewidth, color="black")
    ax.plot((-length/2, length/2), (0, 0), lw=linewidth, color="black")

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,
                                     edgecolor="gray")
    ax.add_patch(p)

def plot_blade_path(ax, R=0.5):
    """Plot blade path as a dashed line."""
    p = plt.Circle((0, 0), R, linestyle="dashed", edgecolor="black",
                   facecolor="none", linewidth=1)
    ax.add_patch(p)

def plot_vectors(fig, ax, theta_deg=0.0, tsr=2.0, c=0.3, label=False):
    """Plot blade velocity, free stream velocity, relative velocity, lift, and
    drag vectors.
    """
    r = 0.5
    u_infty = 0.26
    theta_deg %= 360
    theta_rad = np.deg2rad(theta_deg)
    blade_xy = r*np.cos(theta_rad), r*np.sin(theta_rad)
    head_width = 0.04
    head_length = 0.11
    linewidth = 1.5

    # Function for plotting vector labels
    def plot_label(text, x, y, dx, dy, text_width=0.09, text_height=0.03,
                   sign=-1, dist=1.0/3.0):
        text_width *= plt.rcParams["font.size"]/12*6/fig.get_size_inches()[1]
        text_height *= plt.rcParams["font.size"]/12*6/fig.get_size_inches()[1]
        dvec = np.array((dx, dy))
        perp_vec = rotate(dvec, np.pi/2)
        perp_vec /= mag(perp_vec)
        if theta_deg > 270:
            diag = text_height
        else:
            diag = np.array((text_width, text_height))
        # Projection of text diagonal vector onto normal vector
        proj = np.dot(diag, perp_vec)
        if sign != -1:
            proj = 0 # Text is on right side of vector
        if theta_deg > 180:
            sign *= -1
        dxlab, dylab = perp_vec*(np.abs(proj) + .01)*sign
        xlab, ylab = x + dx*dist + dxlab, y + dy*dist + dylab
        ax.text(xlab, ylab, text)

    # Make blade velocity vector
    x1, y1 = rotate((0.5, tsr*u_infty), np.deg2rad(theta_deg))
    dx, dy = np.array(blade_xy) - np.array((x1, y1))
    blade_vel = np.array((dx, dy))
    ax.arrow(x1, y1, dx, dy, head_width=head_width, head_length=head_length,
             length_includes_head=True, color=dark_gray, linewidth=linewidth)
    if label:
        plot_label(r"$-\omega r$", x1, y1, dx*0.25, dy*0.5)
        # Make chord line vector
        x1c, y1c = np.array((x1, y1)) - np.array((dx, dy))*0.5
        x2c, y2c = np.array((x1, y1)) + np.array((dx, dy))*2
        ax.plot([x1c, x2c], [y1c, y2c], marker=None, color="k", linestyle="-.",
                zorder=1)

    # 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)
    u_infty = np.array((0, -u_infty))
    if label:
        dy = -mag(u_infty)
        plot_label(r"$U_\mathrm{in}$", x1, y1, 0, dy, text_width=0.1)

    # Make relative velocity vector
    dx, dy = np.array(blade_xy) - np.array((x1, y1))
    rel_vel = u_infty + blade_vel
    ax.plot((x1, x1 + dx), (y1, y1 + dy), lw=0)
    ax.arrow(x1, y1, dx, dy, head_width=head_width, head_length=head_length,
             length_includes_head=True, color=tan, linewidth=linewidth)
    if label:
        plot_label(r"$U_\mathrm{rel}$", x1, y1, dx, dy, sign=1,
                   text_width=0.11)

    # Calculate angle between blade vel and rel vel
    alpha_deg = np.rad2deg(np.arccos(np.dot(blade_vel/mag(blade_vel),
                                            rel_vel/mag(rel_vel))))
    if theta_deg > 180:
        alpha_deg *= -1

    # Make drag vector
    drag_amplify = 3.0
    data = lookup_foildata(alpha_deg)
    drag = data["cd"]*mag(rel_vel)**2*drag_amplify
    if drag < 0.4/drag_amplify:
        hs = 0.5
    else:
        hs = 1
    dx, dy = drag*np.array((dx, dy))/mag((dx, dy))
    ax.arrow(blade_xy[0], blade_xy[1], dx, dy, head_width=head_width*hs,
             head_length=head_length*hs, length_includes_head=True, color=red,
             linewidth=linewidth)
    if label:
        plot_label(r"$F_d$", blade_xy[0], blade_xy[1], dx, dy, sign=-1,
                   dist=0.66)

    # Make lift vector
    lift_amplify = 1.5
    lift = data["cl"]*mag(rel_vel)**2*lift_amplify
    dx, dy = rotate((dx, dy), -np.pi/2)/mag((dx, dy))*lift
    if np.abs(lift) < 0.4/lift_amplify:
        hs = 0.5
    else:
        hs = 1
    ax.plot((blade_xy[0], blade_xy[0] + dx), (blade_xy[1], blade_xy[1] + dy),
            linewidth=0)
    ax.arrow(blade_xy[0], blade_xy[1], dx, dy, head_width=head_width*hs,
             head_length=head_length*hs, length_includes_head=True,
             color=green, linewidth=linewidth)
    if label:
        plot_label(r"$F_l$", blade_xy[0], blade_xy[1], dx, dy, sign=-1,
                   text_width=0.12, text_height=0.02, dist=0.66)

    # Label radius
    if label:
        plot_label("$r$", 0, 0, blade_xy[0], blade_xy[1], text_width=0.04,
                   text_height=0.04)

    # Label angle of attack
    if label:
        ast = "simple,head_width={},tail_width={},head_length={}".format(
                head_width*8, linewidth/16, head_length*8)
        xy = blade_xy - rel_vel/mag(rel_vel)*0.2
        ax.annotate(r"$\alpha$", xy=xy, xycoords="data",
                    xytext=(37.5, 22.5), textcoords="offset points",
                    arrowprops=dict(arrowstyle=ast,
                                    ec="none",
                                    connectionstyle="arc3,rad=0.1",
                                    color="k"))
        xy = blade_xy - blade_vel/mag(blade_vel)*0.2
        ax.annotate("", xy=xy, xycoords="data",
                    xytext=(-15, -30), textcoords="offset points",
                    arrowprops=dict(arrowstyle=ast,
                                    ec="none",
                                    connectionstyle="arc3,rad=-0.1",
                                    color="k"))

    # Label azimuthal angle
    if label:
        xy = np.array(blade_xy)*0.6
        ast = "simple,head_width={},tail_width={},head_length={}".format(
                head_width*5.5, linewidth/22, head_length*5.5)
        ax.annotate(r"$\theta$", xy=xy, xycoords="data",
                    xytext=(0.28, 0.12), textcoords="data",
                    arrowprops=dict(arrowstyle=ast,
                                    ec="none",
                                    connectionstyle="arc3,rad=0.1",
                                    color="k"))
        ax.annotate("", xy=(0.41, 0), xycoords="data",
                    xytext=(0.333, 0.12), textcoords="data",
                    arrowprops=dict(arrowstyle=ast,
                                    ec="none",
                                    connectionstyle="arc3,rad=-0.1",
                                    color="k"))

    # Label pitching moment
    if label:
        xy = np.array(blade_xy)*1.1 - blade_vel/mag(blade_vel) * c/4
        ast = "simple,head_width={},tail_width={},head_length={}".format(
                head_width*8, linewidth/16, head_length*8)
        ax.annotate(r"", xy=xy, xycoords="data",
                    xytext=(25, -15), textcoords="offset points",
                    arrowprops=dict(arrowstyle=ast,
                                    ec="none",
                                    connectionstyle="arc3,rad=0.6",
                                    color="k"))
        plot_label(r"$M$", xy[0], xy[1], 0.1, 0.1, sign=-1, dist=0.66)

    return {"u_infty": u_infty, "blade_vel": blade_vel, "rel_vel": rel_vel}

def plot_alpha(ax=None, tsr=2.0, theta=None, alpha_ss=None, **kwargs):
    """Plot angle of attack versus azimuthal angle."""
    if theta is not None:
        theta %= 360
    if ax is None:
        fig, ax = plt.subplots()
    df = calc_cft_ctorque(tsr=tsr)
    ax.plot(df.theta, df.alpha_deg, **kwargs)
    ax.set_ylabel(r"$\alpha$ (degrees)")
    ax.set_xlabel(r"$\theta$ (degrees)")
    ax.set_xlim((0, 360))
    ylim = np.round(df.alpha_deg.max() + 5)
    ax.set_ylim((-ylim, ylim))
    if theta is not None:
        f = interp1d(df.theta, df.alpha_deg)
        ax.plot(theta, f(theta), "ok")
    if alpha_ss is not None:
        ax.hlines((alpha_ss, -alpha_ss), 0, 360, linestyles="dashed")

def plot_rel_vel_mag(ax=None, tsr=2.0, theta=None, **kwargs):
    """Plot relative velocity magnitude versus azimuthal angle."""
    if theta is not None:
        theta %= 360
    if ax is None:
        fig, ax = plt.subplots()
    df = calc_cft_ctorque(tsr=tsr)
    ax.plot(df.theta, df.rel_vel_mag, **kwargs)
    ax.set_ylabel(r"$|\vec{U}_\mathrm{rel}|$")
    ax.set_xlabel(r"$\theta$ (degrees)")
    ax.set_xlim((0, 360))
    if theta is not None:
        f = interp1d(df.theta, df.rel_vel_mag)
        ax.plot(theta, f(theta), "ok")

def plot_alpha_relvel_all(tsrs=np.arange(1.5, 6.1, 1.0), save=False):
    """Plot angle of attack and relative velocity magnitude for a list of TSRs.

    Figure will have two subplots in a single row.
    """
    fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(7.5, 3.0))
    cm = plt.cm.get_cmap("Reds")
    for tsr in tsrs:
        color = cm(tsr/np.max(tsrs))
        plot_alpha(ax=ax1, tsr=tsr, label=r"$\lambda = {}$".format(tsr),
                   color=color)
        plot_rel_vel_mag(ax=ax2, tsr=tsr, color=color)
    [a.set_xticks(np.arange(0, 361, 60)) for a in (ax1, ax2)]
    ax1.legend(loc=(0.17, 1.1), ncol=len(tsrs))
    ax1.set_ylim((-45, 45))
    ax1.set_yticks(np.arange(-45, 46, 15))
    ax2.set_ylabel(r"$|\vec{U}_\mathrm{rel}|/U_\infty$")
    fig.tight_layout()
    if save:
        fig.savefig("figures/alpha_deg_urel_geom.pdf", bbox_inches="tight")

def plot_ctorque(ax=None, tsr=2.0, theta=None, **kwargs):
    """Plot torque coefficient versus azimuthal angle."""
    theta %= 360
    if ax is None:
        fig, ax = plt.subplots()
    df = calc_cft_ctorque(tsr=tsr)
    ax.plot(df.theta, df.ctorque, **kwargs)
    ax.set_ylabel("Torque coeff.")
    ax.set_xlabel(r"$\theta$ (degrees)")
    ax.set_xlim((0, 360))
    if theta is not None:
        f = interp1d(df.theta, df.ctorque)
        ax.plot(theta, f(theta), "ok")

def plot_diagram(fig=None, ax=None, theta_deg=0.0, tsr=2.0, label=False,
                 save=False, axis="on", full_view=True):
    """Plot full vector diagram."""
    if ax is None:
        fig, ax = plt.subplots(figsize=(6, 6))

    plot_blade_path(ax)
    if label:
        # Create dashed line for x-axis
        ax.plot((-0.5, 0.5), (0, 0), linestyle="dashed", color="k",
                zorder=1)
    plot_foil(ax, c=0.3, theta_deg=theta_deg)
    plot_radius(ax, theta_deg)
    plot_center(ax)
    plot_vectors(fig, ax, theta_deg, tsr, label=label)

    # Figure formatting
    if full_view:
        ax.set_xlim((-1, 1))
        ax.set_ylim((-1, 1))
    ax.set_aspect(1)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.axis(axis)

    if save:
        fig.savefig("figures/cft-vectors.pdf")

def plot_all(theta_deg=0.0, tsr=2.0, scale=1.0, full_view=True):
    """Create diagram and plots of kinematics in a single figure."""
    fig = plt.figure(figsize=(7.5*scale, 4.75*scale))
    # Draw vector diagram
    ax1 = plt.subplot2grid((3, 3), (0, 0), colspan=2, rowspan=3)
    plot_diagram(fig, ax1, theta_deg, tsr, axis="on", full_view=full_view)
    # Plot angle of attack
    ax2 = plt.subplot2grid((3, 3), (0, 2))
    plot_alpha(ax2, tsr=tsr, theta=theta_deg, alpha_ss=18, color=light_blue)
    # Plot relative velocity magnitude
    ax3 = plt.subplot2grid((3, 3), (1, 2))
    plot_rel_vel_mag(ax3, tsr=tsr, theta=theta_deg, color=tan)
    # Plot torque coefficient
    ax4 = plt.subplot2grid((3, 3), (2, 2))
    plot_ctorque(ax4, tsr=tsr, theta=theta_deg, color=purple)

    fig.tight_layout()
    return fig

def make_frame(t):
    """Make a frame for a movie."""
    sec_per_rev = 5.0
    deg = t/sec_per_rev*360
    return mplfig_to_npimage(plot_all(deg, scale=2.0))

def make_animation(filetype="mp4", fps=30):
    """Make animation video."""
    if not os.path.isdir("videos"):
        os.mkdir("videos")
    animation = VideoClip(make_frame, duration=5.0)
    if "mp4" in filetype.lower():
        animation.write_videofile("videos/cft-animation.mp4", fps=fps)
    elif "gif" in filetype.lower():
        animation.write_gif("videos/cft-animation.gif", fps=fps)

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Create cross-flow turbine \
                                     vector diagrams.")
    parser.add_argument("create", choices=["figure", "diagram", "animation"],
                        help="Either create a static figure or animation")
    parser.add_argument("--angle", type=float, default=60.0,
                        help="Angle (degrees) to create figure")
    parser.add_argument("--show", action="store_true", default=False)
    parser.add_argument("--save", "-s", action="store_true", default=False,
                        help="Save figure")
    args = parser.parse_args()

    if args.save:
        if not os.path.isdir("figures"):
            os.mkdir("figures")

    if args.create == "diagram":
        set_sns(font_scale=2)
        plot_diagram(theta_deg=args.angle, label=True, axis="off",
                     save=args.save)
    elif args.create == "figure":
        set_sns()
        plot_alpha_relvel_all(save=args.save)
    elif args.create == "animation":
        set_sns(font_scale=2)
        from moviepy.editor import VideoClip
        from moviepy.video.io.bindings import mplfig_to_npimage
        make_animation()

    if args.show:
        plt.show()


"""
==========
Libsvm GUI
==========

A simple graphical frontend for Libsvm mainly intended for didactic
purposes. You can create data points by point and click and visualize
the decision region induced by different kernels and parameter settings.

To create positive examples click the left mouse button; to create
negative examples click the right button.

If all examples are from the same class, it uses a one-class SVM.

"""
from __future__ import division, print_function

print(__doc__)

# Author: Peter Prettenhoer <peter.prettenhofer@gmail.com>
#
# License: BSD 3 clause

import matplotlib
matplotlib.use('TkAgg')

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.backends.backend_tkagg import NavigationToolbar2TkAgg
from matplotlib.figure import Figure
from matplotlib.contour import ContourSet

import Tkinter as Tk
import sys
import numpy as np

from sklearn import svm
from sklearn.datasets import dump_svmlight_file
from sklearn.externals.six.moves import xrange

y_min, y_max = -50, 50
x_min, x_max = -50, 50

class Model(object):
    """The Model which hold the data. It implements the
    observable in the observer pattern and notifies the
    registered observers on change event.
    """

    def __init__(self):
        self.observers = []
        self.surface = None
        self.data = []
        self.cls = None
        self.surface_type = 0

    def changed(self, event):
        """Notify the observers. """
        for observer in self.observers:
            observer.update(event, self)

    def add_observer(self, observer):
        """Register an observer. """
        self.observers.append(observer)

    def set_surface(self, surface):
        self.surface = surface

    def dump_svmlight_file(self, file):
        data = np.array(self.data)
        X = data[:, 0:2]
        y = data[:, 2]
        dump_svmlight_file(X, y, file)

class Controller(object):
    def __init__(self, model):
        self.model = model
        self.kernel = Tk.IntVar()
        self.surface_type = Tk.IntVar()
        # Whether or not a model has been fitted
        self.fitted = False

    def fit(self):
        print("fit the model")
        train = np.array(self.model.data)
        X = train[:, 0:2]
        y = train[:, 2]

        C = float(self.complexity.get())
        gamma = float(self.gamma.get())
        coef0 = float(self.coef0.get())
        degree = int(self.degree.get())
        kernel_map = {0: "linear", 1: "rbf", 2: "poly"}
        if len(np.unique(y)) == 1:
            clf = svm.OneClassSVM(kernel=kernel_map[self.kernel.get()],
                                  gamma=gamma, coef0=coef0, degree=degree)
            clf.fit(X)
        else:
            clf = svm.SVC(kernel=kernel_map[self.kernel.get()], C=C,
                          gamma=gamma, coef0=coef0, degree=degree)
            clf.fit(X, y)
        if hasattr(clf, 'score'):
            print("Accuracy:", clf.score(X, y) * 100)
        X1, X2, Z = self.decision_surface(clf)
        self.model.clf = clf
        self.model.set_surface((X1, X2, Z))
        self.model.surface_type = self.surface_type.get()
        self.fitted = True
        self.model.changed("surface")

    def decision_surface(self, cls):
        delta = 1
        x = np.arange(x_min, x_max + delta, delta)
        y = np.arange(y_min, y_max + delta, delta)
        X1, X2 = np.meshgrid(x, y)
        Z = cls.decision_function(np.c_[X1.ravel(), X2.ravel()])
        Z = Z.reshape(X1.shape)
        return X1, X2, Z

    def clear_data(self):
        self.model.data = []
        self.fitted = False
        self.model.changed("clear")

    def add_example(self, x, y, label):
        self.model.data.append((x, y, label))
        self.model.changed("example_added")

        # update decision surface if already fitted.
        self.refit()

    def refit(self):
        """Refit the model if already fitted. """
        if self.fitted:
            self.fit()

class View(object):
    """Test docstring. """
    def __init__(self, root, controller):
        f = Figure()
        ax = f.add_subplot(111)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_xlim((x_min, x_max))
        ax.set_ylim((y_min, y_max))
        canvas = FigureCanvasTkAgg(f, master=root)
        canvas.show()
        canvas.get_tk_widget().pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
        canvas._tkcanvas.pack(side=Tk.TOP, fill=Tk.BOTH, expand=1)
        canvas.mpl_connect('button_press_event', self.onclick)
        toolbar = NavigationToolbar2TkAgg(canvas, root)
        toolbar.update()
        self.controllbar = ControllBar(root, controller)
        self.f = f
        self.ax = ax
        self.canvas = canvas
        self.controller = controller
        self.contours = []
        self.c_labels = None
        self.plot_kernels()

    def plot_kernels(self):
        self.ax.text(-50, -60, "Linear: $u^T v$")
        self.ax.text(-20, -60, "RBF: $\exp (-\gamma \| u-v \|^2)$")
        self.ax.text(10, -60, "Poly: $(\gamma \, u^T v + r)^d$")

    def onclick(self, event):
        if event.xdata and event.ydata:
            if event.button == 1:
                self.controller.add_example(event.xdata, event.ydata, 1)
            elif event.button == 3:
                self.controller.add_example(event.xdata, event.ydata, -1)

    def update_example(self, model, idx):
        x, y, l = model.data[idx]
        if l == 1:
            color = 'w'
        elif l == -1:
            color = 'k'
        self.ax.plot([x], [y], "%so" % color, scalex=0.0, scaley=0.0)

    def update(self, event, model):
        if event == "examples_loaded":
            for i in xrange(len(model.data)):
                self.update_example(model, i)

        if event == "example_added":
            self.update_example(model, -1)

        if event == "clear":
            self.ax.clear()
            self.ax.set_xticks([])
            self.ax.set_yticks([])
            self.contours = []
            self.c_labels = None
            self.plot_kernels()

        if event == "surface":
            self.remove_surface()
            self.plot_support_vectors(model.clf.support_vectors_)
            self.plot_decision_surface(model.surface, model.surface_type)

        self.canvas.draw()

    def remove_surface(self):
        """Remove old decision surface."""
        if len(self.contours) > 0:
            for contour in self.contours:
                if isinstance(contour, ContourSet):
                    for lineset in contour.collections:
                        lineset.remove()
                else:
                    contour.remove()
            self.contours = []

    def plot_support_vectors(self, support_vectors):
        """Plot the support vectors by placing circles over the
        corresponding data points and adds the circle collection
        to the contours list."""
        cs = self.ax.scatter(support_vectors[:, 0], support_vectors[:, 1],
                             s=80, edgecolors="k", facecolors="none")
        self.contours.append(cs)

    def plot_decision_surface(self, surface, type):
        X1, X2, Z = surface
        if type == 0:
            levels = [-1.0, 0.0, 1.0]
            linestyles = ['dashed', 'solid', 'dashed']
            colors = 'k'
            self.contours.append(self.ax.contour(X1, X2, Z, levels,
                                                 colors=colors,
                                                 linestyles=linestyles))
        elif type == 1:
            self.contours.append(self.ax.contourf(X1, X2, Z, 10,
                                                  cmap=matplotlib.cm.bone,
                                                  origin='lower', alpha=0.85))
            self.contours.append(self.ax.contour(X1, X2, Z, [0.0], colors='k',
                                                 linestyles=['solid']))
        else:
            raise ValueError("surface type unknown")

class ControllBar(object):
    def __init__(self, root, controller):
        fm = Tk.Frame(root)
        kernel_group = Tk.Frame(fm)
        Tk.Radiobutton(kernel_group, text="Linear", variable=controller.kernel,
                       value=0, command=controller.refit).pack(anchor=Tk.W)
        Tk.Radiobutton(kernel_group, text="RBF", variable=controller.kernel,
                       value=1, command=controller.refit).pack(anchor=Tk.W)
        Tk.Radiobutton(kernel_group, text="Poly", variable=controller.kernel,
                       value=2, command=controller.refit).pack(anchor=Tk.W)
        kernel_group.pack(side=Tk.LEFT)

        valbox = Tk.Frame(fm)
        controller.complexity = Tk.StringVar()
        controller.complexity.set("1.0")
        c = Tk.Frame(valbox)
        Tk.Label(c, text="C:", anchor="e", width=7).pack(side=Tk.LEFT)
        Tk.Entry(c, width=6, textvariable=controller.complexity).pack(
            side=Tk.LEFT)
        c.pack()

        controller.gamma = Tk.StringVar()
        controller.gamma.set("0.01")
        g = Tk.Frame(valbox)
        Tk.Label(g, text="gamma:", anchor="e", width=7).pack(side=Tk.LEFT)
        Tk.Entry(g, width=6, textvariable=controller.gamma).pack(side=Tk.LEFT)
        g.pack()

        controller.degree = Tk.StringVar()
        controller.degree.set("3")
        d = Tk.Frame(valbox)
        Tk.Label(d, text="degree:", anchor="e", width=7).pack(side=Tk.LEFT)
        Tk.Entry(d, width=6, textvariable=controller.degree).pack(side=Tk.LEFT)
        d.pack()

        controller.coef0 = Tk.StringVar()
        controller.coef0.set("0")
        r = Tk.Frame(valbox)
        Tk.Label(r, text="coef0:", anchor="e", width=7).pack(side=Tk.LEFT)
        Tk.Entry(r, width=6, textvariable=controller.coef0).pack(side=Tk.LEFT)
        r.pack()
        valbox.pack(side=Tk.LEFT)

        cmap_group = Tk.Frame(fm)
        Tk.Radiobutton(cmap_group, text="Hyperplanes",
                       variable=controller.surface_type, value=0,
                       command=controller.refit).pack(anchor=Tk.W)
        Tk.Radiobutton(cmap_group, text="Surface",
                       variable=controller.surface_type, value=1,
                       command=controller.refit).pack(anchor=Tk.W)

        cmap_group.pack(side=Tk.LEFT)

        train_button = Tk.Button(fm, text='Fit', width=5,
                                 command=controller.fit)
        train_button.pack()
        fm.pack(side=Tk.LEFT)
        Tk.Button(fm, text='Clear', width=5,
                  command=controller.clear_data).pack(side=Tk.LEFT)

def get_parser():
    from optparse import OptionParser
    op = OptionParser()
    op.add_option("--output",
                  action="store", type="str", dest="output",
                  help="Path where to dump data.")
    return op

def main(argv):
    op = get_parser()
    opts, args = op.parse_args(argv[1:])
    root = Tk.Tk()
    model = Model()
    controller = Controller(model)
    root.wm_title("Scikit-learn Libsvm GUI")
    view = View(root, controller)
    model.add_observer(view)
    Tk.mainloop()

    if opts.output:
        model.dump_svmlight_file(opts.output)

if __name__ == "__main__":
    main(sys.argv)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# wqio documentation build configuration file, created by
# sphinx-quickstart on Sun May 22 14:36:00 2016.
#
# This file is execfile()d with the current directory set to its
# containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

import sys
import os
import shlex
import sphinx
import matplotlib as mpl

mpl.use("agg")
import seaborn

clean_bkgd = {"axes.facecolor": "none", "figure.facecolor": "none"}
seaborn.set(style="ticks", rc=clean_bkgd)

numpydoc_show_class_members = False
autodoc_member_order = "bysource"

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
sys.path.insert(0, os.path.abspath("sphinxext"))
extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.doctest",
    "sphinx.ext.intersphinx",
    "sphinx.ext.todo",
    "sphinx.ext.mathjax",
    "sphinx.ext.viewcode",
    #'plot_generator',
    #'plot_directive',
    "numpydoc",
    "ipython_directive",
    "ipython_console_highlighting",
    "sphinx_gallery.gen_gallery",
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# The suffix(es) of source filenames.
# You can specify multiple suffix as a list of string:
# source_suffix = ['.rst', '.md']
source_suffix = ".rst"

# The encoding of source files.
# source_encoding = 'utf-8-sig'

# Include the example source for plots in API docs
plot_include_source = True
plot_formats = [("png", 90)]
plot_html_show_formats = False
plot_html_show_source_link = False

# The encoding of source files.
# source_encoding = 'utf-8-sig'

# The master toctree document.
master_doc = "index"

# General information about the project.
project = "wqio"
copyright = "2016, Paul Hobson (Geosyntec Consultants)"
author = "Paul Hobson (Geosyntec Consultants)"

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# The short X.Y version.
version = "0.5.1"
# The full version, including alpha/beta/rc tags.
release = "0.5.1"

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#
# This is also used if you do content translation via gettext catalogs.
# Usually you set "language" from the command line for these cases.
language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
# today = ''
# Else, today_fmt is used as the format for a strftime call.
# today_fmt = '%B %d, %Y'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This patterns also effect to html_static_path and html_extra_path
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]

# The reST default role (used for this markup: `text`) to use for all
# documents.
# default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
# add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
# add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
# show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = "sphinx"

# A list of ignored prefixes for module index sorting.
# modindex_common_prefix = []

# If true, keep warnings as "system message" paragraphs in the built documents.
# keep_warnings = False

# If true, `todo` and `todoList` produce output, else they produce nothing.
todo_include_todos = False

# -- Options for HTML output ----------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = "sphinx_rtd_theme"

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
# html_theme_options = {}

# Add any paths that contain custom themes here, relative to this directory.
# html_theme_path = []

# The name for this set of Sphinx documents.
# "<project> v<release> documentation" by default.
html_title = 'wqio v0.5.1'

# A shorter title for the navigation bar.  Default is the same as html_title.
# html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
# html_logo = None

# The name of an image file (relative to this directory) to use as a favicon of
# the docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
# html_favicon = None

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

# Add any extra paths that contain custom files (such as robots.txt or
# .htaccess) here, relative to this directory. These files are copied
# directly to the root of the documentation.
# html_extra_path = []

# If not None, a 'Last updated on:' timestamp is inserted at every page
# bottom, using the given strftime format.
# The empty string is equivalent to '%b %d, %Y'.
# html_last_updated_fmt = None

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
# html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
# html_sidebars = {}

# Additional templates that should be rendered to pages, maps page names to
# template names.
# html_additional_pages = {}

# If false, no module index is generated.
# html_domain_indices = True

# If false, no index is generated.
# html_use_index = True

# If true, the index is split into individual pages for each letter.
# html_split_index = False

# If true, links to the reST sources are added to the pages.
# html_show_sourcelink = True

# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
# html_show_sphinx = True

# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
# html_show_copyright = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
# html_use_opensearch = ''

# This is the file name suffix for HTML files (e.g. ".xhtml").
# html_file_suffix = None

# Language to be used for generating the HTML full-text search index.
# Sphinx supports the following languages:
#   'da', 'de', 'en', 'es', 'fi', 'fr', 'h', 'it', 'ja'
#   'nl', 'no', 'pt', 'ro', 'r', 'sv', 'tr', 'zh'
# html_search_language = 'en'

# A dictionary with options for the search language support, empty by default.
# 'ja' uses this config value.
# 'zh' user can custom change `jieba` dictionary path.
# html_search_options = {'type': 'default'}

# The name of a javascript file (relative to the configuration directory) that
# implements a search results scorer. If empty, the default will be used.
# html_search_scorer = 'scorer.js'

# Output file base name for HTML help builder.
htmlhelp_basename = "wqiodoc"

# -- Options for LaTeX output ---------------------------------------------

latex_elements = {
    # The paper size ('letterpaper' or 'a4paper').
    #'papersize': 'letterpaper',
    # The font size ('10pt', '11pt' or '12pt').
    #'pointsize': '10pt',
    # Additional stuff for the LaTeX preamble.
    #'preamble': '',
    # Latex figure (float) alignment
    #'figure_align': 'htbp',
}

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title,
#  author, documentclass [howto, manual, or own class]).
latex_documents = [
    (
        master_doc,
        "wqio.tex",
        "wqio Documentation",
        "Paul Hobson (Geosyntec Consultants)",
        "manual",
    )
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
# latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
# latex_use_parts = False

# If true, show page references after internal links.
# latex_show_pagerefs = False

# If true, show URL addresses after external links.
# latex_show_urls = False

# Documents to append as an appendix to all manuals.
# latex_appendices = []

# If false, no module index is generated.
# latex_domain_indices = True

# -- Options for manual page output ---------------------------------------

# One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section).
man_pages = [(master_doc, "wqio", "wqio Documentation", [author], 1)]

# If true, show URL addresses after external links.
# man_show_urls = False

# -- Options for Texinfo output -------------------------------------------

# Grouping the document tree into Texinfo files. List of tuples
# (source start file, target name, title, author,
#  dir menu entry, description, category)
texinfo_documents = [
    (
        master_doc,
        "wqio",
        "wqio Documentation",
        author,
        "wqio",
        "One line description of project.",
        "Miscellaneous",
    )
]

# Documents to append as an appendix to all manuals.
# texinfo_appendices = []

# If false, no module index is generated.
# texinfo_domain_indices = True

# How to display URL addresses: 'footnote', 'no', or 'inline'.
# texinfo_show_urls = 'footnote'

# If true, do not generate a @detailmenu in the "Top" node's menu.
# texinfo_no_detailmenu = False

# Example configuration for intersphinx: refer to the Python standard library.
intersphinx_mapping = {"https://docs.python.org/": None}


#------------------------------------------------------------------------------
# Name:        pychrono example
# Purpose:
#
# Author:      Alessandro Tasora
#
# Created:     1/01/2019
# Copyright:   (c) ProjectChrono 2019
#------------------------------------------------------------------------------

import pychrono.core as chrono
import pychrono.irrlicht as chronoirr
import matplotlib.pyplot as plt
import numpy as np

print ("Example: create a slider crank and plot results");

# Change this path to asset path, if running from other working dir.
# It must point to the data folder, containing GUI assets (textures, fonts, meshes, etc.)
chrono.SetChronoDataPath("../../../data/")

# ---------------------------------------------------------------------
#
#  Create the simulation system and add items
#

mysystem      = chrono.ChSystemNSC()

# Some data shared in the following
crank_center = chrono.ChVectorD(-1,0.5,0)
crank_rad    = 0.4
crank_thick  = 0.1
rod_length   = 1.5

# Create four rigid bodies: the truss, the crank, the rod, the piston.

# Create the floor truss
mfloor = chrono.ChBodyEasyBox(3, 1, 3, 1000)
mfloor.SetPos(chrono.ChVectorD(0,-0.5,0))
mfloor.SetBodyFixed(True)
mysystem.Add(mfloor)

# Create the flywheel crank
mcrank = chrono.ChBodyEasyCylinder(crank_rad, crank_thick, 1000)
mcrank.SetPos(crank_center + chrono.ChVectorD(0, 0, -0.1))
# Since ChBodyEasyCylinder creates a vertical (y up) cylinder, here rotate it:
mcrank.SetRot(chrono.Q_ROTATE_Y_TO_Z)
mysystem.Add(mcrank)

# Create a stylized rod
mrod = chrono.ChBodyEasyBox(rod_length, 0.1, 0.1, 1000)
mrod.SetPos(crank_center + chrono.ChVectorD(crank_rad+rod_length/2 , 0, 0))
mysystem.Add(mrod)

# Create a stylized piston
mpiston = chrono.ChBodyEasyCylinder(0.2, 0.3, 1000)
mpiston.SetPos(crank_center + chrono.ChVectorD(crank_rad+rod_length, 0, 0))
mpiston.SetRot(chrono.Q_ROTATE_Y_TO_X)
mysystem.Add(mpiston)

# Now create constraints and motors between the bodies.

# Create crank-truss joint: a motor that spins the crank flywheel
my_motor = chrono.ChLinkMotorRotationSpeed()
my_motor.Initialize(mcrank,   # the first connected body
                    mfloor,   # the second connected body
                    chrono.ChFrameD(crank_center)) # where to create the motor in abs.space
my_angularspeed = chrono.ChFunction_Const(chrono.CH_C_PI) # ang.speed: 180°/s
my_motor.SetMotorFunction(my_angularspeed)
mysystem.Add(my_motor)

# Create crank-rod joint
mjointA = chrono.ChLinkLockRevolute()
mjointA.Initialize(mrod,
                   mcrank,
                   chrono.ChCoordsysD( crank_center + chrono.ChVectorD(crank_rad,0,0) ))
mysystem.Add(mjointA)

# Create rod-piston joint
mjointB = chrono.ChLinkLockRevolute()
mjointB.Initialize(mpiston,
                   mrod,
                   chrono.ChCoordsysD( crank_center + chrono.ChVectorD(crank_rad+rod_length,0,0) ))
mysystem.Add(mjointB)

# Create piston-truss joint
mjointC = chrono.ChLinkLockPrismatic()
mjointC.Initialize(mpiston,
                   mfloor,
                   chrono.ChCoordsysD(
                               crank_center + chrono.ChVectorD(crank_rad+rod_length,0,0),
                               chrono.Q_ROTATE_Z_TO_X)
                  )
mysystem.Add(mjointC)

# ---------------------------------------------------------------------
#
#  Create an Irrlicht application to visualize the system
#

myapplication = chronoirr.ChIrrApp(mysystem, 'PyChrono example', chronoirr.dimension2du(1024,768))

myapplication.AddTypicalSky()
myapplication.AddTypicalLogo(chrono.GetChronoDataPath() + 'logo_pychrono_alpha.png')
myapplication.AddTypicalCamera(chronoirr.vector3df(1,1,3), chronoirr.vector3df(0,1,0))
myapplication.AddTypicalLights()

            # ==IMPORTANT!== Use this function for adding a ChIrrNodeAsset to all items
            # in the system. These ChIrrNodeAsset assets are 'proxies' to the Irrlicht meshes.
            # If you need a finer control on which item really needs a visualization proxy in
            # Irrlicht, just use application.AssetBind(myitem); on a per-item basis.

myapplication.AssetBindAll();

            # ==IMPORTANT!== Use this function for 'converting' into Irrlicht meshes the assets
            # that you added to the bodies into 3D shapes, they can be visualized by Irrlicht!

myapplication.AssetUpdateAll();

# ---------------------------------------------------------------------
#
#  Run the simulation
#

# Initialize these lists to store values to plot.
array_time   = []
array_angle  = []
array_pos    = []
array_speed  = []

myapplication.SetTimestep(0.005)

# Run the interactive simulation loop
while(myapplication.GetDevice().run()):

    # for plotting, append instantaneous values:
    array_time.append(mysystem.GetChTime())
    array_angle.append(my_motor.GetMotorRot())
    array_pos.append(mpiston.GetPos().x)
    array_speed.append(mpiston.GetPos_dt().x)

    # here happens the visualization and step time integration
    myapplication.BeginScene()
    myapplication.DrawAll()
    myapplication.DoStep()
    myapplication.EndScene()

    # stop simulation after 2 seconds
    if mysystem.GetChTime() > 2:
          myapplication.GetDevice().closeDevice()

# Use matplotlib to make two plots when simulation ended:
fig, (ax1, ax2) = plt.subplots(2, sharex = True)

ax1.plot(array_angle, array_pos)
ax1.set(ylabel='position [m]')
ax1.grid()

ax2.plot(array_angle, array_speed, 'r--')
ax2.set(ylabel='speed [m]',xlabel='angle [rad]')
ax2.grid()

# trick to plot \pi on x axis of plots instead of 1 2 3 4 etc.
plt.xticks(np.linspace(0, 2*np.pi, 5),['0','$\pi/2$','$\pi$','$3\pi/2$','$2\pi$'])


"""
The :mod:`sklearn.feature_extraction.image` submodule gathers utilities to
extract features from images.
"""

# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
#          Olivier Grisel
#          Vlad Niculae
# License: BSD 3 clause

from itertools import product
import numbers
import numpy as np
from scipy import sparse
from numpy.lib.stride_tricks import as_strided

from ..utils import check_array, check_random_state
from ..base import BaseEstimator

__all__ = ['PatchExtractor',
           'extract_patches_2d',
           'grid_to_graph',
           'img_to_graph',
           'reconstruct_from_patches_2d']

###############################################################################
# From an image to a graph

def _make_edges_3d(n_x, n_y, n_z=1):
    """Returns a list of edges for a 3D image.

    Parameters
    ===========
    n_x : integer
        The size of the grid in the x direction.
    n_y : integer
        The size of the grid in the y direction.
    n_z : integer, optional
        The size of the grid in the z direction, defaults to 1
    """
    vertices = np.arange(n_x * n_y * n_z).reshape((n_x, n_y, n_z))
    edges_deep = np.vstack((vertices[:, :, :-1].ravel(),
                            vertices[:, :, 1:].ravel()))
    edges_right = np.vstack((vertices[:, :-1].ravel(),
                             vertices[:, 1:].ravel()))
    edges_down = np.vstack((vertices[:-1].ravel(), vertices[1:].ravel()))
    edges = np.hstack((edges_deep, edges_right, edges_down))
    return edges

def _compute_gradient_3d(edges, img):
    n_x, n_y, n_z = img.shape
    gradient = np.abs(img[edges[0] // (n_y * n_z),
                      (edges[0] % (n_y * n_z)) // n_z,
                      (edges[0] % (n_y * n_z)) % n_z] -
                      img[edges[1] // (n_y * n_z),
                      (edges[1] % (n_y * n_z)) // n_z,
                      (edges[1] % (n_y * n_z)) % n_z])
    return gradient

# XXX: Why mask the image after computing the weights?

def _mask_edges_weights(mask, edges, weights=None):
    """Apply a mask to edges (weighted or not)"""
    inds = np.arange(mask.size)
    inds = inds[mask.ravel()]
    ind_mask = np.logical_and(np.in1d(edges[0], inds),
                              np.in1d(edges[1], inds))
    edges = edges[:, ind_mask]
    if weights is not None:
        weights = weights[ind_mask]
    if len(edges.ravel()):
        maxval = edges.max()
    else:
        maxval = 0
    order = np.searchsorted(np.unique(edges.ravel()), np.arange(maxval + 1))
    edges = order[edges]
    if weights is None:
        return edges
    else:
        return edges, weights

def _to_graph(n_x, n_y, n_z, mask=None, img=None,
              return_as=sparse.coo_matrix, dtype=None):
    """Auxiliary function for img_to_graph and grid_to_graph
    """
    edges = _make_edges_3d(n_x, n_y, n_z)

    if dtype is None:
        if img is None:
            dtype = np.int
        else:
            dtype = img.dtype

    if img is not None:
        img = np.atleast_3d(img)
        weights = _compute_gradient_3d(edges, img)
        if mask is not None:
            edges, weights = _mask_edges_weights(mask, edges, weights)
            diag = img.squeeze()[mask]
        else:
            diag = img.ravel()
        n_voxels = diag.size
    else:
        if mask is not None:
            mask = mask.astype(dtype=np.bool, copy=False)
            mask = np.asarray(mask, dtype=np.bool)
            edges = _mask_edges_weights(mask, edges)
            n_voxels = np.sum(mask)
        else:
            n_voxels = n_x * n_y * n_z
        weights = np.ones(edges.shape[1], dtype=dtype)
        diag = np.ones(n_voxels, dtype=dtype)

    diag_idx = np.arange(n_voxels)
    i_idx = np.hstack((edges[0], edges[1]))
    j_idx = np.hstack((edges[1], edges[0]))
    graph = sparse.coo_matrix((np.hstack((weights, weights, diag)),
                              (np.hstack((i_idx, diag_idx)),
                               np.hstack((j_idx, diag_idx)))),
                              (n_voxels, n_voxels),
                              dtype=dtype)
    if return_as is np.ndarray:
        return graph.toarray()
    return return_as(graph)

def img_to_graph(img, mask=None, return_as=sparse.coo_matrix, dtype=None):
    """Graph of the pixel-to-pixel gradient connections

    Edges are weighted with the gradient values.

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    img : ndarray, 2D or 3D
        2D or 3D image
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as : np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype : None or dtype, optional
        The data of the returned sparse matrix. By default it is the
        dtype of img

    Notes
    -----
    For scikit-learn versions 0.14.1 and prior, return_as=np.ndarray was
    handled by returning a dense np.matrix instance.  Going forward, np.ndarray
    returns an np.ndarray, as expected.

    For compatibility, user code relying on this method should wrap its
    calls in ``np.asarray`` to avoid type issues.
    """
    img = np.atleast_3d(img)
    n_x, n_y, n_z = img.shape
    return _to_graph(n_x, n_y, n_z, mask, img, return_as, dtype)

def grid_to_graph(n_x, n_y, n_z=1, mask=None, return_as=sparse.coo_matrix,
                  dtype=np.int):
    """Graph of the pixel-to-pixel connections

    Edges exist if 2 voxels are connected.

    Parameters
    ----------
    n_x : int
        Dimension in x axis
    n_y : int
        Dimension in y axis
    n_z : int, optional, default 1
        Dimension in z axis
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as : np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype : dtype, optional, default int
        The data of the returned sparse matrix. By default it is int

    Notes
    -----
    For scikit-learn versions 0.14.1 and prior, return_as=np.ndarray was
    handled by returning a dense np.matrix instance.  Going forward, np.ndarray
    returns an np.ndarray, as expected.

    For compatibility, user code relying on this method should wrap its
    calls in ``np.asarray`` to avoid type issues.
    """
    return _to_graph(n_x, n_y, n_z, mask=mask, return_as=return_as,
                     dtype=dtype)

###############################################################################
# From an image to a set of small image patches

def _compute_n_patches(i_h, i_w, p_h, p_w, max_patches=None):
    """Compute the number of patches that will be extracted in an image.

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    i_h : int
        The image height
    i_w : int
        The image with
    p_h : int
        The height of a patch
    p_w : int
        The width of a patch
    max_patches : integer or float, optional default is None
        The maximum number of patches to extract. If max_patches is a float
        between 0 and 1, it is taken to be a proportion of the total number
        of patches.
    """
    n_h = i_h - p_h + 1
    n_w = i_w - p_w + 1
    all_patches = n_h * n_w

    if max_patches:
        if (isinstance(max_patches, (numbers.Integral))
                and max_patches < all_patches):
            return max_patches
        elif (isinstance(max_patches, (numbers.Real))
                and 0 < max_patches < 1):
            return int(max_patches * all_patches)
        else:
            raise ValueError("Invalid value for max_patches: %r" % max_patches)
    else:
        return all_patches

def extract_patches(arr, patch_shape=8, extraction_step=1):
    """Extracts patches of any n-dimensional array in place using strides.

    Given an n-dimensional array it will return a 2n-dimensional array with
    the first n dimensions indexing patch position and the last n indexing
    the patch content. This operation is immediate (O(1)). A reshape
    performed on the first n dimensions will cause numpy to copy data, leading
    to a list of extracted patches.

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    arr : ndarray
        n-dimensional array of which patches are to be extracted

    patch_shape : integer or tuple of length arr.ndim
        Indicates the shape of the patches to be extracted. If an
        integer is given, the shape will be a hypercube of
        sidelength given by its value.

    extraction_step : integer or tuple of length arr.ndim
        Indicates step size at which extraction shall be performed.
        If integer is given, then the step is uniform in all dimensions.

    Returns
    -------
    patches : strided ndarray
        2n-dimensional array indexing patches on first n dimensions and
        containing patches on the last n dimensions. These dimensions
        are fake, but this way no data is copied. A simple reshape invokes
        a copying operation to obtain a list of patches:
        result.reshape([-1] + list(patch_shape))
    """

    arr_ndim = arr.ndim

    if isinstance(patch_shape, numbers.Number):
        patch_shape = tuple([patch_shape] * arr_ndim)
    if isinstance(extraction_step, numbers.Number):
        extraction_step = tuple([extraction_step] * arr_ndim)

    patch_strides = arr.strides

    slices = [slice(None, None, st) for st in extraction_step]
    indexing_strides = arr[slices].strides

    patch_indices_shape = ((np.array(arr.shape) - np.array(patch_shape)) //
                           np.array(extraction_step)) + 1

    shape = tuple(list(patch_indices_shape) + list(patch_shape))
    strides = tuple(list(indexing_strides) + list(patch_strides))

    patches = as_strided(arr, shape=shape, strides=strides)
    return patches

def extract_patches_2d(image, patch_size, max_patches=None, random_state=None):
    """Reshape a 2D image into a collection of patches

    The resulting patches are allocated in a dedicated array.

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    image : array, shape = (image_height, image_width) or
        (image_height, image_width, n_channels)
        The original image data. For color images, the last dimension specifies
        the channel: a RGB image would have `n_channels=3`.

    patch_size : tuple of ints (patch_height, patch_width)
        the dimensions of one patch

    max_patches : integer or float, optional default is None
        The maximum number of patches to extract. If max_patches is a float
        between 0 and 1, it is taken to be a proportion of the total number
        of patches.

    random_state : int, RandomState instance or None, optional (default=None)
        Pseudo number generator state used for random sampling to use if
        `max_patches` is not None.  If int, random_state is the seed used by
        the random number generator; If RandomState instance, random_state is
        the random number generator; If None, the random number generator is
        the RandomState instance used by `np.random`.

    Returns
    -------
    patches : array, shape = (n_patches, patch_height, patch_width) or
         (n_patches, patch_height, patch_width, n_channels)
         The collection of patches extracted from the image, where `n_patches`
         is either `max_patches` or the total number of patches that can be
         extracted.

    Examples
    --------

    >>> from sklearn.feature_extraction import image
    >>> one_image = np.arange(16).reshape((4, 4))
    >>> one_image
    array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11],
           [12, 13, 14, 15]])
    >>> patches = image.extract_patches_2d(one_image, (2, 2))
    >>> print(patches.shape)
    (9, 2, 2)
    >>> patches[0]
    array([[0, 1],
           [4, 5]])
    >>> patches[1]
    array([[1, 2],
           [5, 6]])
    >>> patches[8]
    array([[10, 11],
           [14, 15]])
    """
    i_h, i_w = image.shape[:2]
    p_h, p_w = patch_size

    if p_h > i_h:
        raise ValueError("Height of the patch should be less than the height"
                         " of the image.")

    if p_w > i_w:
        raise ValueError("Width of the patch should be less than the width"
                         " of the image.")

    image = check_array(image, allow_nd=True)
    image = image.reshape((i_h, i_w, -1))
    n_colors = image.shape[-1]

    extracted_patches = extract_patches(image,
                                        patch_shape=(p_h, p_w, n_colors),
                                        extraction_step=1)

    n_patches = _compute_n_patches(i_h, i_w, p_h, p_w, max_patches)
    if max_patches:
        rng = check_random_state(random_state)
        i_s = rng.randint(i_h - p_h + 1, size=n_patches)
        j_s = rng.randint(i_w - p_w + 1, size=n_patches)
        patches = extracted_patches[i_s, j_s, 0]
    else:
        patches = extracted_patches

    patches = patches.reshape(-1, p_h, p_w, n_colors)
    # remove the color dimension if useless
    if patches.shape[-1] == 1:
        return patches.reshape((n_patches, p_h, p_w))
    else:
        return patches

def reconstruct_from_patches_2d(patches, image_size):
    """Reconstruct the image from all of its patches.

    Patches are assumed to overlap and the image is constructed by filling in
    the patches from left to right, top to bottom, averaging the overlapping
    regions.

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    patches : array, shape = (n_patches, patch_height, patch_width) or
        (n_patches, patch_height, patch_width, n_channels)
        The complete set of patches. If the patches contain colour information,
        channels are indexed along the last dimension: RGB patches would
        have `n_channels=3`.

    image_size : tuple of ints (image_height, image_width) or
        (image_height, image_width, n_channels)
        the size of the image that will be reconstructed

    Returns
    -------
    image : array, shape = image_size
        the reconstructed image

    """
    i_h, i_w = image_size[:2]
    p_h, p_w = patches.shape[1:3]
    img = np.zeros(image_size)
    # compute the dimensions of the patches array
    n_h = i_h - p_h + 1
    n_w = i_w - p_w + 1
    for p, (i, j) in zip(patches, product(range(n_h), range(n_w))):
        img[i:i + p_h, j:j + p_w] += p

    for i in range(i_h):
        for j in range(i_w):
            # divide by the amount of overlap
            # XXX: is this the most efficient way? memory-wise yes, cpu wise?
            img[i, j] /= float(min(i + 1, p_h, i_h - i) *
                               min(j + 1, p_w, i_w - j))
    return img

class PatchExtractor(BaseEstimator):
    """Extracts patches from a collection of images

    Read more in the :ref:`User Guide <image_feature_extraction>`.

    Parameters
    ----------
    patch_size : tuple of ints (patch_height, patch_width)
        the dimensions of one patch

    max_patches : integer or float, optional default is None
        The maximum number of patches per image to extract. If max_patches is a
        float in (0, 1), it is taken to mean a proportion of the total number
        of patches.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    """
    def __init__(self, patch_size=None, max_patches=None, random_state=None):
        self.patch_size = patch_size
        self.max_patches = max_patches
        self.random_state = random_state

    def fit(self, X, y=None):
        """Do nothing and return the estimator unchanged

        This method is just there to implement the usual API and hence
        work in pipelines.
        """
        return self

    def transform(self, X):
        """Transforms the image samples in X into a matrix of patch data.

        Parameters
        ----------
        X : array, shape = (n_samples, image_height, image_width) or
            (n_samples, image_height, image_width, n_channels)
            Array of images from which to extract patches. For color images,
            the last dimension specifies the channel: a RGB image would have
            `n_channels=3`.

        Returns
        -------
        patches : array, shape = (n_patches, patch_height, patch_width) or
             (n_patches, patch_height, patch_width, n_channels)
             The collection of patches extracted from the images, where
             `n_patches` is either `n_samples * max_patches` or the total
             number of patches that can be extracted.

        """
        self.random_state = check_random_state(self.random_state)
        n_images, i_h, i_w = X.shape[:3]
        X = np.reshape(X, (n_images, i_h, i_w, -1))
        n_channels = X.shape[-1]
        if self.patch_size is None:
            patch_size = i_h // 10, i_w // 10
        else:
            patch_size = self.patch_size

        # compute the dimensions of the patches array
        p_h, p_w = patch_size
        n_patches = _compute_n_patches(i_h, i_w, p_h, p_w, self.max_patches)
        patches_shape = (n_images * n_patches,) + patch_size
        if n_channels > 1:
            patches_shape += (n_channels,)

        # extract the patches
        patches = np.empty(patches_shape)
        for ii, image in enumerate(X):
            patches[ii * n_patches:(ii + 1) * n_patches] = extract_patches_2d(
                image, patch_size, self.max_patches, self.random_state)
        return patches