import numpy
from model import *
from squaremesh import *
from setmask import *
from parameterize import *
from setflowequation import *
from EnumDefinitions import *
from solve import *
from MatlabFuncs import *
from PythonFuncs import *

"""
This test is a test from the ISMP-HOM Intercomparison project.
Pattyn and Payne 2006
"""

L_list=[5000.,10000.,20000.,40000.,80000.,160000.]
results=[]

for L in L_list:
	nx=20    #numberof nodes in x direction
	ny=20
	md=model()
	md=squaremesh(md,L,L,nx,ny)
	md=setmask(md,'','')    #ice sheet test
	md=parameterize(md,'../Par/ISMIPB.py')
	md.extrude(10,1.)
	md=setflowequation(md,'pattyn','all')

	#Create dirichlet on the bed only
	md.diagnostic.spcvx=float('NaN')*numpy.ones((md.mesh.numberofvertices,1))
	md.diagnostic.spcvy=float('NaN')*numpy.ones((md.mesh.numberofvertices,1))
	md.diagnostic.spcvz=float('NaN')*numpy.ones((md.mesh.numberofvertices,1))

	pos=numpy.nonzero(md.mesh.vertexonbed)
	md.diagnostic.spcvx[pos]=0.
	md.diagnostic.spcvy[pos]=0.

	#Create MPCs to have periodic boundary conditions
	posx=numpy.nonzero(md.mesh.x==0.)[0]
	posx2=numpy.nonzero(md.mesh.x==numpy.max(md.mesh.x))[0]

	posy=numpy.nonzero(logical_and_n(md.mesh.y==0.,md.mesh.x!=0.,md.mesh.x!=numpy.max(md.mesh.x)))[0]    #Don't take the same nodes two times
	posy2=numpy.nonzero(logical_and_n(md.mesh.y==numpy.max(md.mesh.y),md.mesh.x!=0.,md.mesh.x!=numpy.max(md.mesh.x)))[0]

	md.diagnostic.vertex_pairing=numpy.vstack((numpy.hstack((posx.reshape(-1,1)+1,posx2.reshape(-1,1)+1)),numpy.hstack((posy.reshape(-1,1)+1,posy2.reshape(-1,1)+1))))

	#Compute the diagnostic
	md.diagnostic.abstol=float('NaN')
	md.cluster=generic('name',oshostname(),'np',8)
	md=solve(md,DiagnosticSolutionEnum())
	pos=numpy.nonzero(logical_or_n(md.mesh.x==0.,md.mesh.y==0.,md.mesh.x==numpy.max(md.mesh.x),md.mesh.y==numpy.max(md.mesh.y)))
	md.diagnostic.spcvx[pos]=md.results.DiagnosticSolution.Vx[pos]
	md.diagnostic.spcvy[pos]=md.results.DiagnosticSolution.Vy[pos]
	md.diagnostic.vertex_pairing=numpy.empty((0,2),int)
	md=setflowequation(md,'stokes','all')
	md=solve(md,DiagnosticSolutionEnum())

	#Plot the results and save them
	vx=md.results.DiagnosticSolution.Vx
	vy=md.results.DiagnosticSolution.Vy
	vz=md.results.DiagnosticSolution.Vz
	results.append(md.results.DiagnosticSolution)

#	plotmodel(md,'data',vx,'data',vy,'data',vz,'layer#all',md.mesh.numberoflayers)

#Fields and tolerances to track changes
field_names     =[\
	'Vx5km','Vy5km','Vz5km',\
	'Vx10km','Vy10km','Vz10km',\
	'Vx20km','Vy20km','Vz20km',\
	'Vx40km','Vy40km','Vz40km',\
	'Vx80km','Vy80km','Vz80km',\
	'Vx160km','Vy160km','Vz160km'
]
field_tolerances=[\
	1e-08,1e-08,1e-08,\
	1e-08,1e-08,1e-08,\
	1e-08,1e-08,1e-08,\
	1e-08,1e-08,1e-08,\
	1e-08,1e-07,1e-08,\
	1e-07,1e-07,1e-07,\
]
field_values=[]
for result in results:
	field_values=field_values+[\
		result.Vx,\
		result.Vy,\
		result.Vz,\
		]
