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

last_surface_editing_2d_system = None
def surface_editing_2d_system_cached( undeformed, h ):
    '''
    Same as surface_editing_2d_system(), but caches the system for successive identical calls.
    '''
    
    global last_surface_editing_2d_system
    
    N,M = undeformed.shape
    
    if not last_surface_editing_2d_system or (N,M,h) != last_surface_editing_2d_system[0]:
        last_surface_editing_2d_system = ((N,M,h), surface_editing_2d_system( undeformed, h ))
    
    return copy( last_surface_editing_2d_system[1][0] ), copy( last_surface_editing_2d_system[1][1] )

def ind2ij( N,M, ind ):
    assert ind >= 0 and ind < N*M
    return ind // M, ind % M
def ij2ind( N,M, i,j ):
    assert i >= 0 and i < N and j >= 0 and j <= M
    return i*M + j


def surface_editing_2d_system( undeformed, h ):
    '''
    Returns a system that for solving a surface_editing_2d problem given
    undeformed positions 'undeformed',
    and distance between elements 'h'.
    '''
    
    N,M = undeformed.shape
    
    # function computing the system matrix
    def compute_hessian_direct():
        H_direct = zerosf( (N*M,N*M) )
        coeff = 1. / (h*h)
        
        # evaluate element A[row,col] (ignoring boundary conditions)
        def element( row,col ):
            i0,j0 = ind2ij( N,M, row )
            i1,j1 = ind2ij( N,M, col )
            if (i0,j0) == (i1,j1): return -4
            elif i0 == i1 and abs(j0 - j1) == 1: return 1
            elif j0 == j1 and abs(i0 - i1) == 1: return 1
            else: return 0.
    
        for i in xrange( N*M ):
            for j in xrange( N*M ):
                H_direct[ i,j ] = element( i,j )
        
        return (coeff*coeff)*dot( H_direct.T, H_direct )
    
    A = compute_hessian_direct()
    #print A
    # right-hand side of the system, initially zeros
    b = zerosf(N*M)

    # setting boundary conditions

    # conditions on two rows of boundary indices, for positions and derivatives
    for ind in xrange( N*M ):
        i,j = ind2ij( N,M, ind )
        if i == 0 or i == 1 or j == 0 or j == 1 or i == N-2 or i == N-1 or j == M-2 or j == M-1:
            ## set rhs values to undeformed values
            b[ind] = undeformed[i,j]
            ## set corresponding rows of the matrix to identity rows
            A[ind,:] = identityf( shape(A)[0] )[ind,:]
    
    return A,b

def surface_editing_2d( 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'.
    '''
    
    N,M = undeformed.shape
    A,b = surface_editing_2d_system_cached( undeformed, h )
    
    # for indices for which values are fixed, set
    # matrix rows to identity rows, and rhs entries to the fixed values
    I = identityf( shape(A)[0] )
    for (i,j), val in fixed.iteritems():
        ind = ij2ind( N,M, i,j )
        A[ind,:] = I[ind,:]
        b[ind] = val
    
    # solve the system
    #print A
    Ainv = inv(A)
    
    x = dot( Ainv, b )
    #x = Ainv *b 
    
    return x.reshape( (N,M) )


def test1():
    col_steps = 7
    row_steps = 5
    undeformed = zerosf((row_steps,col_steps))
    h = 1
    fixed = { (row_steps//2, col_steps//2): 1.0 }
    deformed = surface_editing_2d( fixed, undeformed, h )
    print deformed[row_steps//2,:]

def main():
    test1()

if __name__ == '__main__': main()