import pylab
import matplotlib.axes3d as axes3d
from numpy import *
from myarray import *
from direct_linear_2d import surface_editing_2d

def get_coord_val( event ):
    assert len( fixed ) == 1
    
    if event.button != 2: return None
    
    ## get_coord() doesn't work
    #return get_coord( event ), fixed.values()[0]
    
    return fixed.keys()[0], get_value( event )

def get_value( event ):
    return float( event.y ) / ax3d.bbox.height()

def get_coord( event ):
    '''
    For some reason, event.xdata and event.ydata are wrong, so I can't use these.
    '''
    colcoord = int( round( event.xdata ) )
    rowcoord = int( round( event.ydata ) )
    
    if colcoord < 0 or colcoord >= num_cols: return None
    if rowcoord < 0 or rowcoord >= num_rows: return None
    
    return (colcoord, rowcoord)


# on mouse button down, enable mouse motion callback
# and update the surface
def on_press(event):
    global motion_id
    
    print( 'press', event.xdata, event.ydata )
    
    coord_val = get_coord_val( event )
    if coord_val:
        # enable callback for dragging
        motion_id = pylab.connect('motion_notify_event', on_motion)
        
        update_fixed( *coord_val )


# on mouse button release, desable motion callback
def on_release(event):
    global motion_id
    pylab.disconnect(motion_id)

# mouse motion callback: update fixed point position,
# recompute the surface and redraw the plot
def on_motion(event):
    print( 'motion', event.xdata, event.ydata )
    
    coord_val = get_coord_val( event )
    if coord_val:
        update_fixed( *coord_val )

def on_key(event):
    '''
    Handles key press events.
    '''
    
    assert len( fixed ) == 1
    N = num_rows
    M = num_cols
    ( row,col ), val = fixed.items()[0]
    
    if event.key == 'left':
        col = max( 2, min( M-3, col - 1 ) )
    elif event.key == 'right':
        col = max( 2, min( M-3, col + 1 ) )
    elif event.key == 'down':
        row = max( 2, min( N-3, row - 1 ) )
    elif event.key == 'up':
        row = max( 2, min( N-3, row + 1 ) )
    else:
        return
    
    update_fixed( ( row, col ), val )


def update_fixed( coord, val ):
    global fixed, deformed, ax3d
    
    # change fixed point
    fixed = { coord: val }
    print fixed
    
    # recompute the surface, and update surface data in the plot object
    deformed = surface_editing_2d( fixed, undeformed, h )
    ax3d.cla()
    #surf = ax3d.plot_surface(xs, ys, deformed)
    surf = ax3d.plot_wireframe(xs, ys, deformed)
    surf.set_array(linspace(0, 1.0, 100))
    
    ax3d.set_xlim( 0, num_cols )
    ax3d.set_ylim( 0, num_rows )
    ax3d.set_zlim( 0, 1 )
    
    ax3d.set_xlabel( "X" )
    ax3d.set_ylabel( "Y" )
    ax3d.set_zlabel( "Z" )
    
    pylab.draw()


# number of sample points on the curve
num_cols = 20
num_rows = 20
# initial positions of curve points: all zeros
# in this version of the code, only the first and the last 2 are used
undeformed = zerosf((num_rows,num_cols))
# step
h = 1
# intial position for a fixed point in the center
fixed = {(num_rows//2,num_cols//2): 1.0}
# solve for positions satisfying this constraint
deformed = surface_editing_2d(fixed, undeformed, h)

fig = pylab.gcf()
ax3d = axes3d.Axes3D(fig)
plt = fig.axes.append(ax3d)

xs,ys = meshgrid(range(num_cols),range(num_rows))

update_fixed( fixed.keys()[0], fixed.values()[0] )

# initially, no callback for mouse motion
motion_id = 0
# callback for mouse button press and release
press_id = pylab.connect('button_press_event', on_press)
release_id = pylab.connect('button_release_event', on_release)
keypress_id = pylab.connect('key_press_event', on_key)
pylab.show()