from pylab import *
from myarray import *
import sys
import time

import springs_dynamic

def on_press(event):
    '''
    On mouse button down, if a point is close to the mouse, make it a constraint and drag the point.
    '''
    global motion_id, pts, down_pos, drag_ind, down_was_constraint
    prox_threshold = 0.05
    
    assert len(pts) > 0
    
    down_pos = ( event.xdata, event.ydata )
    if down_pos[0] is None or down_pos[1] is None: return
    
    # index of closest point on the list
    closest_ind = argmin(map( lambda z: norm( z - down_pos), pts ))
    drag_ind = -1
    
    # if there is a point close to the click point, drag it
    if norm( pts[closest_ind] - down_pos ) < prox_threshold:
        # enable callback for dragging
        motion_id = connect('motion_notify_event', on_motion)
        drag_ind = closest_ind
        
        pts[ drag_ind ] = down_pos
        down_was_constraint = drag_ind in constraints
        if not down_was_constraint: constraints.append( drag_ind )
    
    update_points_and_edges()

def on_release(event):
    '''
    On mouse button release, disable motion callback.
    If the point was clicked and not dragged, toggle its constrained-ness.
    '''
    global motion_id, pts, drag_ind, down_pos, constraints, iterate_duration_after_mouse_up
    click_threshold = 0.05
    
    if -1 != drag_ind and down_was_constraint and norm( pts[ drag_ind ] - down_pos ) < click_threshold:
        constraints.remove( drag_ind )
    
    drag_ind = -1
    disconnect(motion_id)
    
    if manual_iteration:
        update_points_and_edges()
    else:
        update_points_and_edges_for_duration( iterate_duration_after_mouse_up )

def on_motion(event):
    '''
    On mouse drag, update the dragged point's constraint position, recompute the spring,
    and redraw the plot.
    '''
    global pts, drag_ind
    
    ## We shouldn't see this, but I don't want to assert because window systems do the
    ## darndest things.
    if drag_ind != -1:
        
        new_pt = ( event.xdata, event.ydata )
        if new_pt[0] is None or new_pt[1] is None: return
        
        pts[ drag_ind ] = new_pt
    
    update_points_and_edges()

def on_key(event):
    '''
    Key press handler.
    '''
    global manual_iteration
    
    if event.key == 'e': set_explicit( True )
    elif event.key == 'i': set_explicit( False )
    elif event.key == ' ': update_points_and_edges()
    elif event.key == 'm':
        manual_iteration = not manual_iteration
        print 'manual iteration:', manual_iteration

def update_points_and_edges():
    global pts, vs
    
    ## Compute new locations for the points.
    pts, vs = stepper( pts, vs, edges, edges_rest_lengths, constraints, dt, verbose = False )
    
    ## Tell matplotlib about the new point positions.
    pts_un = pts.tolist()
    pts_c = []
    
    constraint_indices = list( constraints )
    constraint_indices.sort()
    constraint_indices.reverse()
    for c in constraint_indices:
        del pts_un[ c ]
        pts_c.append( pts[ c ] )
    
    pts_un = asarrayf( pts_un )
    pts_c = asarrayf( pts_c )
    
    plt_un.set_data(pts.T[0],pts.T[1])
    if len( pts_c ): plt_c.set_data(pts_c.T[0],pts_c.T[1])
    else: plt_c.set_data([],[])
    
    
    ## Tell matplotlib about the new edge positions.
    for i, edge in enumerate( edges ):
        pt0 = pts[ edge[0] ]
        pt1 = pts[ edge[1] ]
        plt_edges[i].set_data( [ pt0[0], pt1[0] ], [ pt0[1], pt1[1] ] )
    
    draw()

def update_points_and_edges_for_duration( seconds ):
    '''
    Calls update_points_and_edges() in a loop until 'seconds' of cpu time have elapsed.
    '''
    processor_time_end = time.clock() + seconds
    while time.clock() < processor_time_end:
        update_points_and_edges()


pts = None
edges = None

def make_simple_chain():
    global pts, edges, constraints
    
    ## A super-simple chain of three nodes
    pts = arrayf( [ (-.5,0), (0,0), (.5,0) ] )
    edges = [ (0,1), (1,2) ]
    
    constraints = [ 0, 1 ]

def make_grid( n, m ):
    global pts, edges, constraints
    
    assert n > 1 and m > 1
    
    ## I want to make a grid of n-by-m points in the range [-.5, .5] x [-.5, .5].
    ## So a point i,j 
    
    origin = arrayf( (-.5, -.5) )
    nstep = 1. / (n - 1)
    mstep = 1. / (m - 1)
    
    pts = []
    edges = []
    for i in xrange( n ):
        for j in xrange( m ):
            pts.append( origin + (j*mstep, i*nstep) )
            
            if i > 0: edges.append( ( i*m + j, (i-1)*m + (j-0) ) )
            if j > 0: edges.append( ( i*m + j, (i-0)*m + (j-1) ) )
            if j > 0 and i > 0:
                edges.append( ( (i-0)*m + j, (i-1)*m + (j-1) ) )
                edges.append( ( (i-1)*m + j, (i-0)*m + (j-1) ) )
    
    pts = arrayf( pts )
    
    constraints = [ 0, m-1 ]

#make_simple_chain()
#make_grid( 2, 3 )
make_grid( 4, 4 )

vs = 0.05 * pts

#edges_rest_lengths = springs_dynamic.compute_edge_lengths( 0 * pts, edges )
edges_rest_lengths = springs_dynamic.compute_edge_lengths( pts, edges )

## The default is implicit (backward Euler) integration.
## If 'explicit' is specified on the command line, then explicit (forward Euler) integration will
## be used.
def set_explicit( do_explicit ):
    global dt, nsteps_per_draw, stepper
    
    if do_explicit:
        dt = 0.1
        stepper = springs_dynamic.forward_euler
        print 'Using forward euler integration.'
    else:
        dt = 1
        stepper = springs_dynamic.backward_euler
        print 'Using backward euler integration.'

set_explicit( 'explicit' in sys.argv[1:] )
manual_iteration = False
iterate_duration_after_mouse_up = 40

drag_ind = -1
down_pos = (0,0)
down_was_constraint = True

plt_edges = [ plot([],'k-')[0] for e in edges ]
plt_un = plot([],'bo')[0]
plt_c = plot([],'ro')[0]

axis([-1,1,-1,1])

## Attach callbacks for mouse button press and release.
press_id = connect('button_press_event', on_press)
release_id = connect('button_release_event', on_release)
# Initially, no callback for mouse motion.
motion_id = 0
## Uncomment this line to make the motion callback (wiggling the mouse) update the spring system
# motion_id = connect('motion_notify_event', on_motion)

## Attach callback for keyboard input.
key_id = connect('key_release_event', on_key)

update_points_and_edges()

show()