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 curve(rfun,tmin=-2,tmax=2,tpts=200,color='black'):
    domain = np.linspace(tmin,tmax,tpts)
    r = [[rfun(t)[i] for t in domain] for i in range(3)]
    trace = Scatter3d(x=r[0],y=r[1],z=r[2],mode='lines',
                      line=Line(color=color,width=3))
    return(trace)

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,tmin=-2,tmax=2,tpts=20,umin=-2,umax=2,upts=20,color='black'):
    t, u = np.linspace(tmin,tmax,tpts), np.linspace(umin,umax,upts)
    path = snake(t,u)
    r = [[rfun(t,u)[i] for (t,u) in path] for i in range(3)]
    trace = Scatter3d(x=r[0],y=r[1],z=r[2],mode='lines',
                      line=Line(color=color,width=3))
    return(trace)

f = lambda x,y: 1-x*x-x*y-y*y
x0=1
y0=2
xmin=-3
xmax=3
ymin=-3
ymax=3
s = surface(lambda t,u: (t,u,f(t,u)),color='green',tmin=xmin,tmax=xmax,umin=ymin,umax=ymax) #surface z=f(x,y)
sy = curve(lambda t: (t,y0,f(t,y0)),color='red',tmin=xmin,tmax=xmax) #slice at y=y0
sx = curve(lambda t: (x0,t,f(x0,t)),color='blue',tmin=ymin,tmax=ymax) #slice at x=x0
py.iplot(Figure(data=Data([s,sy,sx])))

f = lambda x,y: sin(5*x*y)
x0=1
y0=2
xmin=-3
xmax=3
ymin=-3
ymax=3
s = surface(lambda t,u: (t,u,f(t,u)),color='green',
            tmin=xmin,tmax=xmax,umin=ymin,umax=ymax,tpts=80,upts=80) #surface z=f(x,y)
sy = curve(lambda t: (t,y0,f(t,y0)),color='red',tmin=xmin,tmax=xmax) #slice at y=y0
sx = curve(lambda t: (x0,t,f(x0,t)),color='blue',tmin=ymin,tmax=ymax) #slice at x=x0
py.iplot(Figure(data=Data([s,sy,sx])))