#!/usr/bin/env python
#asdf test
from __future__ import division
from numpy import *
from koko.fab.asdf import ASDF
from koko.c.vec3f import Vec3f
from koko.c.mat3f import Mat3f
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import csv
import sys
import time

def lame(E,nu):
	mu = E/(2*(1+nu))
	l = E*nu/(1+nu)/(1-2*nu)
	return mu,l
def print_bounds(asdf):
	print "ASDF: (%f,%f),(%f,%f),(%f,%f)"%(asdf.xmin,asdf.xmax,asdf.ymin,asdf.ymax,asdf.zmin,asdf.zmax)

print 'opening body asdf'
body = ASDF.load('boundary/body.asdf')
print 'opened body asdf of depth %d and %d cells'%(body.depth, body.cell_count)
print_bounds(body)

print 'opening boundary condition asdf'
boundary = ASDF.load('boundary/support.asdf')
print 'opened boundary condition asdf of depth %d and %d cells'%(boundary.depth, boundary.cell_count)
print_bounds(body)

degree = 1
#Remember, ADSF units are mm!
F = Vec3f(0,0,-10.*1e0) #gravity, N =  m/s2
E = 1e5/(1e6) #youngs modulus, Pa = N/m2, normalized to N/mm
nu = .3
mu,l = lame(E,nu)
h = 1.25 #for this equally spaced version, this parameter determines number of shape functions

#create grid of points for shape function zeros
Xi = rollaxis(mgrid[body.xmin:body.xmax:h,
										body.ymin:body.ymax:h,
										body.zmin:body.zmax:h],0,4)
sh = shape(Xi)
#Xi = Xi.reshape(-1,3)
#n = shape(Xi)[0]

assemble = 0
if assemble:
	print "%d shape functions (including all three axes)"%(3*n)
	A = zeros((3*n,3*n)) #stiffness matrix
	b = zeros(3*n) #rhs
	nn = n
	t0 = time.time()
	for i in range(nn):
		v = body.integrate_bodyforce(boundary,degree,F,Vec3f(Xi[i]),h)
		b[3*i:3*i+3] = v.to_numpy()
	t1 = time.time()
	print "RHS: %d of %d took %f s"%(nn,n,t1-t0)
	for i in range(nn):
		for j in range(nn):
			A[3*i:3*i+3, 3*j:3*j+3] =  body.integrate_stiffness(boundary,degree,Vec3f(Xi[i]),Vec3f(Xi[j]),h,mu,l).to_numpy()
	print "Stiffness: %d of %d took %f s"%(nn,n,time.time()-t1)

	outfile = open('boundary/matrices.npz','w')
	savez(outfile, A=A, b=b)
	outfile.close()

else:
	infile = open('boundary/matrices.npz','r')
	from matplotlib.patches import Rectangle
	npzfile = load(infile)
	A = npzfile['A']
	b = npzfile['b']
	print shape(A),shape(b)
	#z_idx = where(all(A==0,axis=0))
	#print z_idx
	t0 = time.time()
	c = linalg.solve(A,b)
	print "Solved linear system in %f s"%(time.time()-t0)
	from matplotlib import pyplot as plt
	plt.figure(figsize=(16,5))

	body = plt.Rectangle((-26.796999, -1.5), 26.796999, 3., facecolor="#eeddee")
	support = plt.Rectangle((0, -2.5), 2.5, 5, facecolor="#aaeeaa")
	plt.gca().add_artist(body)
	body.set_alpha(.3)
	plt.gca().add_artist(support)
	support.set_alpha(.3)

	c = c.reshape(sh)
	colors = array([[191,15,15],[255,50,25],[255,125,38],[164,126,41],[0,90,94]])/256.
	#print shape(c)
	for i in range(5):
		#Q = plt.quiver(c[:,i,:,0].T,c[:,i,:,2].T,color=colors[i],pivot='mid')
		Q = plt.quiver(Xi[:,:,i,0],Xi[:,:,i,1],c[:,i,:,0],c[:,i,:,2],color=colors[i],pivot='mid')
	l,r,b,t = plt.axis()
	dx, dy = r-l, t-b
	plt.axis([l-0.05*dx, r+0.05*dx, b-0.2*dy, t+0.2*dy])
	plt.gca().set_aspect(1.)



	#plt.show()
	#plt.plot(c)
	#plt.show()

	#cn = linalg.norm(c,axis=-1)
	#cz_mid = c[:,2,:,2]
	#print shape(cz_mid)
	#for zslice in range(shape(cz_mid)[-1]):
	#	plt.plot(cz_mid[:,zslice],label='slice:%d'%zslice)
plt.show()