from pylab import *
from myarray import *

# on mouse button down, if a point is close to mouse, enable dragging the point
# otherwise add a point
def on_press(event):
    global motion_id,  pts, drag_ind
    prox_threshold = 0.05
    # index of closest point on the list
    if len(pts): closest_ind = argmin(map( lambda z: norm( z - [event.xdata,event.ydata]), pts[:-1]))

    # if there are no points, create two and we're done.
    # (We need two because of our illustrator-like assumption that every click creates a
    # point and two more to sets its tangent for the bezier segments on either side of it.)
    if len(pts) == 0:
        new_pt = (event.xdata, event.ydata)
        pts = arrayf( [new_pt, new_pt] )
    # if there is a point close to the click point, drag it
    elif len(pts) and norm( pts[closest_ind] - [event.xdata,event.ydata]) < prox_threshold:
        # enable callback for dragging
        motion_id = connect('motion_notify_event', on_motion)
        pts[closest_ind] = [event.xdata, event.ydata]
        drag_ind = closest_ind
    else:
        # create three new point 
        pts_list = pts.tolist();
        new_pt = (event.xdata, event.ydata)
        pts_list.append( new_pt )
        pts_list.append( new_pt )
        pts_list.append( new_pt )
        pts = arrayf( pts_list )
        # add illustrator dragging callback
        motion_id = connect('motion_notify_event', on_illustrator_motion)
    
    update_points()

# on mouse button release, disable motion callback
def on_release(event):
    global motion_id, plt, plt_bspline
    drag_ind = -1
    disconnect(motion_id)
    update_points_and_beziers()

# mouse motion callback: update fixed point position,
# recompute the curve and redraw the plot
def on_motion(event):
    global drag_ind, pts, plt
    if drag_ind >= 0 and drag_ind < shape(pts)[0]:
        pts[drag_ind] = [event.xdata,event.ydata]
        update_points()

def on_illustrator_motion(event):
    '''
    Moves the last point and third-to-last point about the second-to-last point.
    This has the effect of controlling the tangent between 
    '''
    
    global pts
    
    ## We need 3 points to do anything.
    if len( pts ) < 3: return
    
    new_pt = (event.xdata, event.ydata)
    pts[-1] = new_pt
    pts[-3] = 2 * pts[-2] - pts[-1]
    
    update_points()

def on_key(event):
    global pts,plt
    if event.key == 'd' and len(pts) >= 2:
        pts_list = pts.tolist()
        pts_list.pop()
        pts_list.pop()
        if len( pts_list ) > 0: pts_list.pop()
        pts = array(pts_list)
        update_points_and_beziers()

def eval_bspline( pts, degree ):
    from bezier import Bezier_Curve_n
    bspline = Bezier_Curve_n( pts, 3 )
    
    eval_n = 10 * len(pts) + 1
    
    result = []
    for val in linspace( 0, 1, eval_n ):
        result.append( bspline( val ) )
    
    #print 'pts:'
    #print pts
    #print 'result:'
    #print result
    
    return asarrayf( result )

def update_points():
    assert len( pts ) == 0 or len( pts ) == 2 or (len( pts ) - 2) % 3 == 0
    
    if len( pts ) > 0:
        ## Don't show the last point; it's the dangling tangent for the next bezier segment.
        #plt.set_data(pts.T[0][:-1],pts.T[1][:-1])
        ## On second thought, show the last point so it doesn't appear abruptly when the next
        ## segment is added.
        plt.set_data(pts.T[0],pts.T[1])
    else:
        plt.set_data([],[])
    
    draw()

def update_points_and_beziers():
    assert len( pts ) == 0 or len( pts ) == 2 or (len( pts ) - 2) % 3 == 0
    
    if len( pts ) > 0:
        ## Don't show the last point; it's the dangling tangent for the next bezier segment.
        #plt.set_data(pts.T[0][:-1],pts.T[1][:-1])
        ## On second thought, show the last point so it doesn't appear abruptly when the next
        ## segment is added.
        plt.set_data(pts.T[0],pts.T[1])
    else:
        plt.set_data([],[])
    
    if len(pts) >= degree:
        eval_spline_points = eval_bspline(pts[:-1],degree)
        plt_bspline.set_data( eval_spline_points.T[0], eval_spline_points.T[1] )
    else:
        plt_bspline.set_data([],[])
    
    draw()

drag_ind = -1
pts = arrayf([])
plt = plot([],'ro')[0]
plt_bspline = plot([],'b-')[0]
axis([-1,1,-1,1])

degree = 3

# initially, no callback for mouse motion
motion_id = 0
# callback for mouse button press and release
press_id = connect('button_press_event', on_press)
release_id = connect('button_release_event', on_release)
key_id = connect('key_release_event', on_key)

show()