from pylab import *
from myarray import *
from numpy.linalg import solve


# 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.01
    # index of closest point on the list
    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 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, desable motion callback
def on_release(event):
    global motion_id,plt
    drag_ind = -1
    disconnect(motion_id)
    if len(pts) >= 5:
        c = least_sq_polynom(pts,2)
        clf()
        plot_implicit(c,2)
        plt = plot(pts.T[0],pts.T[1],'o')[0]


# 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 poly_powers(n):
    ''' list of 2-ples representing powers of monomials in a bivariate polynomial of degree n, excluding (0,0)'''
    ppowers = [ (i,j) for i in range(0,n+1) for j in range(0,n-i+1)]
    ppowers.pop(0)
    return ppowers

def least_sq_polynom(pts,n):
    ''' fit an implicit polynomial of degree n to points; may fail for degenerate cases'''
    ppowers = poly_powers(n)
    A = matrix( [ map( lambda pp: x**pp[0]*y**pp[1], ppowers ) for (x,y) in pts])
    print(A)
    ATA = dot(A.T , A )
    y = ones((len(pts),1))
    ATy = dot(A.T,y)
    c =solve(ATA, ATy)
    return ravel(c)

def plot_implicit(c,n):
    ''' plot an implicit polynomial curve of deg n  given by its list of coefficients'''
    global X, Y
    ppowers = poly_powers(n)
    cind = dict( zip( ppowers, range(0,len(ppowers))))
    Z = reshape(
        map( lambda x,y:
             sum(  map( lambda ind: c[cind[ind]]*x**ind[0]*y**ind[1], ppowers
                   )),
             ravel(X), ravel(Y)),shape(X))
    contour(X,Y,Z,[1.0])


drag_ind = -1
pts = 0.9*arrayf( [[1.,0.],[0.,1.],[-1.,0.],[0.,-1.],[sqrt(2.)/2.,sqrt(2.)/2.]])
pts = arrayf([[0.5, 0.0]])
plt = plot(pts.T[0],pts.T[1],'o')[0]
axis([-1,1,-1,1])
poly_deg = 2
poly_coeff = []

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