#!/Users/pwd/Library/Enthought/Canopy_64bit/User/bin/python

from numpy import * #if just import numpy, then have to do "numpy.function" each time
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib import cm

def main(args):

    force = lambda x, y: sin(3*pi*x) #(RHS of equation, ie f)

    domain_x = [0,1]; #array with two entries - left point and right point (looks like an interval, but isn't!)
    domain_y = [0,1]; #array with two entries - left point and right point (looks like an interval, but isn't!)
    
    Nx = 51; #number of nodes to discretise x-domain
    Ny = 51; #number of nodes to discretise y-domain

    coords_x = linspace( domain_x[0], domain_x[1], Nx ) #vector that contains the x-coordinates of the nodes (linearly spaced)
    coords_y = linspace( domain_y[0], domain_y[1], Ny ) #vector that contains the y-coordinates of the nodes (linearly spaced)

    vertices = []; #array of array containing coords for each numbered node (ie vertices[node no.] = [x coord, y coord] )
    for j in range(0, Ny):
        for i in range(0, Nx):
            coords = [ coords_x[i], coords_y[j] ];
            vertices.append(coords);

    elements = []; #array of elements containing node numbers for each element (ie elements[element no.] = [node1, node2, node3] )
    for j in range(0, Ny - 1):
        for i in range(0, Nx - 1):
            triangle_upper = [ (j+1)*Ny + i, j*Ny + i, (j+1)*Ny + i+1 ]; #triangles numbered ACW, with hypotenuse between second two nodes
            triangle_lower = [ j*Ny + i+1, (j+1)*Ny + i+1, j*Ny + i ]; #triangles numbered ACW, with hypotenuse between second two nodes
            elements.append(triangle_upper);
            elements.append(triangle_lower);

    bc = {}; #creates empty dictionary for boundary conditions 
    for i in range(0, Nx): #bottom row
        bc[i] = 0;
    for j in range(1, Ny - 1): #sides
        bc[j*Nx] = 0;
        bc[(j+1)*(Nx) - 1] = 0;
    for i in range(0, Nx): #top row
        bc[Nx*(Ny-1) + i] = 0;

    dof = {}; #creates empty dictionary for degrees of freedom
    j = 0; #counter for degrees of freedom
    for i in range(0, Nx*Ny):
        if not (i in bc): #if i is not in the dictionary of boundaries, ie it's not constrained
            dof[i] = j #it is a degree of freedom - include it here
            j += 1

    Ndof = len(dof);

    K = zeros( [Ndof, Ndof] ); #stiffness matrix

    gradN1 = matrix([[-1],[-1]]); gradN2 = matrix([[1],[0]]); gradN3 = matrix([[0],[1]]); #gradient functions on reference triangle
    ref_gradients = array([ gradN1, gradN2, gradN3 ]); #array containing gradient functions for reference triangle
    one = float(1);
    third = one/3;
    f = zeros( Ndof ); 

    #stiffness matrix K
    for k in elements: #iterate over all elements (triangles)
        x1 = vertices[k[0]][0]; y1 = vertices[k[0]][1]; #retrieve coords for nodes
        x2 = vertices[k[1]][0]; y2 = vertices[k[1]][1]; #retrieve coords for nodes
        x3 = vertices[k[2]][0]; y3 = vertices[k[2]][1]; #retrieve coords for nodes
        B = matrix([ [x2-x1, x3-x1], [y2-y1, y3-y1] ]); #set up transformation matrix
        BTi = linalg.inv(B.transpose()); #calculate transpose inverse of transformation matrix
        size = 0.5*linalg.det(B); #calculate the size of the element
        centre = matrix([[third*(x1+x2+x3)],[third*(y1+y2+y3)]])
        for i in range(0,3): #iterate over pairs of vertices
            if(k[i] in dof):
                f[dof[k[i]]] += force(centre[0],centre[1])*third*size
            for j in range(0,3): #iterate over pairs of vertices
                if ((k[i] in dof) and (k[j] in dof)): #only consider those where i & j are in dof
                    di = dof[k[i]]; dj = dof[k[j]]; #find the corresponding dof no.s
                    gradi_ref = ref_gradients[i]; gradj_ref = ref_gradients[j];
                    gradi = BTi*gradi_ref; gradj = BTi*gradj_ref; #calculate the gradients using reference triangle
                    gradi = squeeze(asarray(gradi)); gradj = squeeze(asarray(gradj)); #turn gradients into array (ie vector) for dot producting
                    integrand = -dot(gradj, gradi); #calculate the integrand
                    integral = integrand*size; #calculate the "integral"
                    K[di,dj] += integral;
                    
    #print "K=";
    #print K;
    #print '\n'.join(map(str, dof)); #print the list dof
    #print '\n'.join(map(str, bc)); #print the list bc

    #f = ones( Ndof );

    print "Now Solving";

    u = linalg.solve(K,f);

    #print "u=";
    #print '\n'.join(map(str, u)); #print the resulting u

    """
    plot_x = [];
    plot_y = [];
    plot_z = [];
    for i in range(0, Nx*Ny):
        plot_x.append(vertices[i][0]);
        plot_y.append(vertices[i][1]);
        if(i in bc):
            plot_z.append(bc[i]);
        if(i in dof):
            plot_z.append(u[dof[i]]);

    print "plot_x=";
    print '\n'.join(map(str, plot_x)); #print the list plot_x
    print "plot_y=";
    print '\n'.join(map(str, plot_y)); #print the list plot_y
    print "plot_z=";
    print '\n'.join(map(str, plot_z)); #print the list plot_z

    #print "X";
    #print '\n'.join(map(str, X)); #print the list X
    #print "Y";
    #print '\n'.join(map(str, Y)); #print the list Y
    #print "Z";
    #print '\n'.join(map(str, Z)); #print the list Z

    """

    print "Now Plotting";

    #PLOT
    plot_x = []; #set up empty coord plot lists
    plot_y = []; #set up empty coord plot lists
    plot_z = []; #set up empty coord plot lists
       
    for j in range(0, Nx*Ny): #iterate over all elements (triangles)
        plot_x.append(vertices[j][0]) #add x-coord of vertex to plot lists
        plot_y.append(vertices[j][1]) #add y-coords of vertex to plot lists
        if j in dof:
            plot_z.append(u[dof[j]]); #add z-coord of vertex to plot lists
        else:
            plot_z.append(bc[j]);
        
    fig = plt.figure()
    ax = fig.gca(projection='3d')

    ax.plot_trisurf(plot_x, plot_y, plot_z, cmap=cm.jet, linewidth=0.2)  

    plt.show()
    
if __name__ == '__main__':
    import sys
    main(sys.argv)