from numpy import *
from myarray import *

from springs import compute_edge_lengths

def forward_euler( p_n, v_n, edges, edge_rest_lengths, constraints, dt, verbose = True ):
    '''
    Given a list of points 'p' as an n-by-2 array, a list of velocities 'v' as an n-by-2 array,
    a list of (i,j) pairs 'edges' denoting an edge between points p[i] and p[j],
    a list of rest lengths (one for each edge in 'edges'),
    and a list of constraints denoting p[i] doesn't change,
    takes one forward Euler step of size 'dt' and returns the p, v arrays corresponding to the new
    time step.
    
    NOTE: 'edges' must not have both (i,j) and (j,i)
    '''
    
    from springs import F
    
    ## f = m*a
    ## We assume m = 1, so multiplying by 1/m has no effect.
    a_n = F( p_n, edges, edge_rest_lengths ).reshape( p_n.shape )
    if verbose:
        print 'a_n:'
        print a_n
    
    ## Apply the explicit update rules.
    ## NOTE: With no damping, this oscillates a lot.  One way to apply damping is simply to
    ##       multiply velocity by a factor < 1.
    damping = 1.
    #damping = .99
    v_n_p_1 = damping * v_n + dt * a_n
    p_n_p_1 =           p_n + dt * v_n
    
    ## Apply the constraints.
    for i in constraints:
        p_n_p_1[i] = p_n[i]
        v_n_p_1[i] = 0.
    
    return p_n_p_1, v_n_p_1

def backward_euler( p_n, v_n, edges, edge_rest_lengths, constraints, dt, verbose = True ):
    '''
    Given a list of points 'p' as an n-by-2 array, a list of velocities 'v' as an n-by-2 array,
    a list of (i,j) pairs 'edges' denoting an edge between points p[i] and p[j],
    a list of rest lengths (one for each edge in 'edges'),
    and a list of constraints denoting p[i] doesn't change,
    takes one backward Euler step of size 'dt' and returns the p, v arrays corresponding to the new
    time step.
    
    NOTE: 'edges' must not have both (i,j) and (j,i)
    '''
    
    from springs import F, J, constrain_system
    
    XSTEP_THRESHOLD = 1e-5
    F_THRESHOLD = 1e-8
    MAX_ITERATIONS = 100
    
    N = prod( p_n.shape )
    dim = p_n.shape[1]
    
    ## Combine positions (p) and velocities (v) into a single, large vector.
    pv_n = emptyf( 2 * N )
    pv_n[ : N ] = p_n.flatten()
    pv_n[ N : ] = v_n.flatten()
    
    ## Apply constraints.
    constrain_rows = []
    for i in constraints:
        ## constrain positions
        constrain_rows.extend( range( i*dim, (i+1) * dim ) )
        ## constraint velocities
        constrain_rows.extend( range( N + i*dim, N + (i+1) * dim ) )
        ## zero initial velocities
        pv_n[ N + i*dim : N + (i+1) * dim ] = 0.
    
    ## Store the current positions and velocities.
    pv_n0 = pv_n.copy()
    
    ## g() is a function of pv (positions and velocities) that is zero when
    ## dt * (the derivative at pv) = pv - pv_n0,
    ## where pv_n0 is our initial pv.
    def g( pv ):
        va = emptyf( 2*N )
        va[ :N ] = pv[ N: ]
        va[ N: ] = F( pv[ :N ].reshape( p_n.shape ), edges, edge_rest_lengths )
        return dt*va - pv + pv_n0
    
    ## dg_dpv() computes the Jacobian of g with respect to pv.
    def dg_dpv( pv ):
        dg = zerosf( (2*N,2*N) )
        dg[ :N, N: ] = identityf( N )
        dg[ N:, :N ] = J( pv[ :N ].reshape( p_n.shape ), edges, edge_rest_lengths )
        return dt * dg - identityf( 2*N )
    
    ## Newton-Raphson iterations to find the solution, a satisfying pv.
    iteration = 0
    while True:
        if verbose: print '-- iteration', iteration, '--'
        iteration += 1
        
        g_n = g( pv_n )
        dg_n = dg_dpv( pv_n )
        mag2_g_n = sum( g_n ** 2 )
        if verbose: print '| g( pv_n ) |^2:', mag2_g_n
        if mag2_g_n < F_THRESHOLD: break
        
        constrain_system( dg_n, g_n, constrain_rows )
        
        # pv_n_p_1 = pv_n - dot( linalg.inv( dg_n ), g_n )
        ## <=> pv_n_p_1 - pv_n = -linalg.inv( dg_n ) * g_n
        ## <=> pv_n - pv_n_p_1 = linalg.inv( dg_n ) * g_n
        ## <=> dg_n * ( pv_n - pv_n_p_1 ) = g_n
        p_negative_delta = linalg.solve( dg_n, g_n )
        ## pv_n - ( pv_n - pv_n_p_1 ) = pv_n_p_1
        pv_n_p_1 = pv_n - p_negative_delta
        
        diff2 = sum( ( pv_n_p_1 - pv_n ) ** 2 )
        if verbose: print '| pv_n+1 - pv_n |^2:', diff2
        pv_n = pv_n_p_1
        if diff2 < XSTEP_THRESHOLD: break
        
        if iteration >= MAX_ITERATIONS:
            print 'Diverged.'
            return p_n.copy(), v_n.copy()
            break
    
    return pv_n[ :N ].reshape( p_n.shape ), pv_n[ N: ].reshape( v_n.shape )

def test1():
    print '===== test1() ====='
    
    #p_undeformed = arrayf( [[0,0], [1,1]] )
    #p_undeformed = arrayf( [[0,0], [1,0], [1,1]] )
    #p_undeformed = arrayf( [[0,0], [1,0], [0,1], [1,1]] )
    p_undeformed = arrayf( [[0,0], [1,0], [2,0]] )
    print 'p.shape:', p_undeformed.shape
    print 'p undeformed:', p_undeformed
    
    #edges = [ (0,1) ]
    edges = [ (0,1), (1,2) ]
    #edges = [ (0,1), (0,2), (1,3), (2,3) ]
    print 'edges:', edges
    ## Multiply p_undeformed by 0 to force 0 rest length springs
    #edge_rest_lengths = compute_edge_lengths( 0 * p_undeformed, edges )
    edge_rest_lengths = compute_edge_lengths( p_undeformed, edges )
    print 'edge_rest_lengths:', edge_rest_lengths
    
    #constraints = [ ( 0, p_undeformed[0] ) ]
    #constraints = [ ( 0, p_undeformed[0] ), ( 3, p_undeformed[3] ) ]
    constraints = [ 0, 2 ]
    print 'constraints:', constraints
    
    p_initial = p_undeformed.copy()
    p_initial[1] += array( (.5,0) )
    print 'p initial:', p_initial
    
    v_initial = 0. * p_initial
    print 'v initial:', v_initial
    
    import sys
    explicit = 'explicit' in sys.argv[1:]
    if explicit:
        dt = 1
        N = 1
        stepper, stepper_name = forward_euler, 'forward euler'
    else:
        dt = 1
        N = 1
        stepper, stepper_name = backward_euler, 'backward euler'
    
    t = 0.
    p_step = p_initial.copy()
    v_step = v_initial.copy()
    for n in range(N):
        p_step, v_step = stepper( p_step, v_step, edges, edge_rest_lengths, constraints, dt )
        t += dt
        print '%s step, t = %s' % (stepper_name, t)
        print 'p:'
        print p_step.round(4)
        print 'v:'
        print v_step.round(4)
    

def main():
    test1()

if __name__ == '__main__': main()