from numpy import *
from numpy.linalg import *
from myarray import *
import sys

def curve_editing_1d_system( undeformed, h ):
    '''
    Returns the system matrix and right-hand-side whose solution is
    new positions given
    undeformed positions 'undeformed',
    and distance between elements 'h'.
    NOTE: The returned system is missing boundary conditions.
    '''

    N = len( undeformed )

    # function computing the system matrix     
    def hessian_direct():
        H_direct = zerosf( (N,N) )

        # evaluate element A[i,j] (ignoring boundary conditions)
        def element( i,j ):
            coeff = 1. / (h*h*h*h)
            if i == j: return 6 * coeff
            elif i+1 == j: return -4 * coeff
            elif i-1 == j: return -4 * coeff
            elif i+2 == j: return 1 * coeff
            elif i-2 == j: return 1 * coeff
            else: return 0.

        for i in range( N ):
            for j in range( N ):
                H_direct[ i,j ] = element( i,j )
        return H_direct
    
    A = hessian_direct()
    # right-hand side of the system, initially zeros
    b = zerosf(N)
    
    return A,b


def curve_editing_1d_system_fixed( fixed, undeformed, h ):
    '''
    Returns the system matrix and right-hand-side whose solution is
    new positions given target fixed position dict 'fixed',
    consisting of pairs (index, value)
    undeformed positions 'undeformed',
    and distance between elements 'h'.
    '''
    
    A,b = curve_editing_1d_system( undeformed, h )
    
    # setting boundary conditions

    # conditions at two endpoints, for positions and derivatives
    b[0] = undeformed[0]
    b[1] = undeformed[1]
    b[-2] = undeformed[-2]
    b[-1] = undeformed[-1]

    #  set corresponding rows of the matrix to identity rows   
    A[0,:] = identityf( shape(A)[0] )[0,:]
    A[1,:] = identityf( shape(A)[0] )[1,:]
    A[-2,:] = identityf( shape(A)[0] )[-2,:]
    A[-1,:] = identityf( shape(A)[0] )[-1,:]
    

    # for indices for which values are fixed, set
    # matrix rows to identity rows, and rhs entries to the fixed values
    for i, val in fixed.iteritems():
        A[i,:] = identityf( shape(A)[0] )[i,:]
        b[i] = val
    
    return A, b



def curve_editing_1d( fixed, undeformed, h ):
    '''
    Returns new positions given target fixed position dict 'fixed',
    consisting of pairs (index, value)
    undeformed positions 'undeformed',
    and distance between elements 'h'.
    '''
    
    A,b = curve_editing_1d_system_fixed( fixed, undeformed, h )
    
    # solve the system   
    Ainv = inv(A)
    
    x = dot( Ainv, b )
    #x = Ainv *b 
    
    return x