import plotly.offline as py
from plotly.graph_objs import *
import numpy as np
from numpy import pi, cos, sin, exp, log, sqrt
py.init_notebook_mode()
def snake(u,v):
    lu, lv = len(u), len(v)
    path = []

    i, j = 0, 0
    istep, jstep = 1, 1
    while(i < lu):
        while(0 <= j < lv):
            path.append((u[i],v[j]))
            j += jstep
        j -= jstep
        i += istep
        jstep *= -1

    i -= istep
    i -= istep
    istep *= -1
    while(0 <= j < lv):
        while(0 <= i < lu):
            path.append((u[i],v[j]))
            i += istep
        i -= istep
        j += jstep
        istep *= -1
    return path
def surface(rfun,amin=-2,amax=2,apts=20,bmin=-2,bmax=2,bpts=20,color='black',name=' '):
    u, v = np.linspace(amin,amax,apts), np.linspace(bmin,bmax,bpts)
    path = snake(u,v)
    r = [[rfun(u,v)[i] for (u,v) in path] for i in range(3)]
    trace = Scatter3d(x=r[0],y=r[1],z=r[2],mode='lines',name=name,
                      line=Line(color=color,width=3))
    return(trace)

from random import random
zero=np.array([0,0,0])
def rotation_matrix(axis, theta):
    #based on http://stackoverflow.com/a/6802723
    """
    Return the rotation matrix associated with counterclockwise rotation about
    the given axis by theta radians.
    """
    axis = np.asarray(axis)
    theta = np.asarray(theta)
    axis = axis/sqrt(np.dot(axis, axis))
    a = cos(theta/2.0)
    b, c, d = -axis*sin(theta/2.0)
    aa, bb, cc, dd = a*a, b*b, c*c, d*d
    bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
    return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
                     [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
                     [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])
def arrow(r0=zero,dr=zero,color='black',head_frac=0.1,head_ang=pi/12,head_count=7,name=' '):
    r0=np.asarray(r0)
    dr=np.asarray(dr)
    a=np.cross(dr,[random() for i in range(3)]) #rotation axis perpendicular to dr
    [ra,rb]=[rotation_matrix(axis,angle) for (axis,angle) in [(a,head_ang),(dr,2*pi/head_count)]]
    piece=-head_frac*np.dot(ra,dr)
    steps=[r0,dr,piece,-piece]
    for i in range(head_count-1):
        piece=np.dot(rb,piece)
        steps += [piece,-piece]
    path,position=[],zero
    for step in steps:
        position = position + step
        path.append(position)
    lp=[list(p) for p in path]
    [x,y,z]=list(zip(*lp))
    trace = Scatter3d(x=x,y=y,z=z,mode='lines',line=Line(color=color,width=3),name=name)
    return(trace)

r = lambda u,v: np.asarray([u**2,u+v,v**2])
ndA = lambda u,v: np.asarray([2*v,-4*u*v,2*u])
def normal(u,v): 
    vec = ndA(u,v)
    return vec/sqrt(np.dot(vec,vec))
F = lambda u,v: np.asarray([0,0,u**2])
v_pts = [.1+n*.15 for n in range(7)]
u_pts =[v-.05 for v in v_pts]
uv_pts = [(u,v) for u in u_pts for v in v_pts if u<v]
uf = lambda t,v: t*v
s = surface(lambda t,v: tuple(r(uf(t,v),v)),
                amin=0,bmin=0,amax=1,bmax=1,color='yellow',name='S')
normals=[arrow(r0=r(*p),dr=normal(*p)*.3,color='red',name='n') for p in uv_pts]
field=[arrow(r0=r(*p),dr=F(*p),color='blue',name='F') for p in uv_pts]
py.iplot(Figure(data=Data([s]+normals+field)))