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))

    # if there is a point close to the click point, drag it
    if 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 a new point 
       pts_list = pts.tolist();
       pts_list.append([event.xdata, event.ydata])
       pts = arrayf( pts_list)
    plt.set_data(pts.T[0],pts.T[1])
    draw()

# on mouse button release, disable motion callback
def on_release(event):
    global motion_id, plt, plt_bspline
    drag_ind = -1
    disconnect(motion_id)
    if len(pts) >= degree:
        eval_spline_points = eval_bspline(pts,degree)
        plt_bspline.set_data( eval_spline_points.T[0], eval_spline_points.T[1] )
        plt.set_data( pts.T[0],pts.T[1] )
        draw()


# 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]
        plt.set_data(pts.T[0],pts.T[1])
        draw()

def on_key(event):
    global pts,plt
    if event.key == 'd' and len(pts) > 1:
        pts_list = pts.tolist()
        pts_list.pop()
        pts = array(pts_list)
        plt.set_data(pts.T[0],pts.T[1])
        draw()

def eval_bspline( pts, degree ):
    '''
    Evaluated the spline with points 'pts' and 'degree', using the global spline_gen function
    to get the evaluator.
    '''
    
    global spline_gen
    bspline = spline_gen( pts, degree )
    
    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 eval_spline_gen( type ):
    '''
    Returns a function object that can be called with a list of points
    and a requested spline degree.
    The returned function object then returns an evaluator for the spline, taking a parameter 't'.
    '''
    
    if type == 'bspline':
        from bspline import Uniform_B_Spline
        return lambda P, d: Uniform_B_Spline( P, d )
    elif type == 'cubic_bspline':
        from bspline import Uniform_Cubic_B_Spline
        return lambda P, d: Uniform_Cubic_B_Spline( P )
    elif type == 'bezier':
        from bezier import Bezier_Curve
        return lambda P, d: Bezier_Curve( P )
    else: # type == 'cubic_bezier':
        if not ( type == 'cubic_bezier' ):
            import sys
            print >> sys.stderr, 'WARNING: Unknown spline type: %s.  Defaulting to cubic_bezier.' % (type,)
        from bezier import Bezier_Curve_n
        return lambda P, d: Bezier_Curve_n( P, d )


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

degree = 3
spline_type = 'cubic_bezier'
import sys
if len( sys.argv ) == 2: spline_type = sys.argv[1]
spline_gen = eval_spline_gen( spline_type )

[X,Y] = meshgrid( arange(-1,1,0.1), arange(-1,1,0.1))

# 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()