import numpy
from model import *
from triangle import *
from setmask import *
from parameterize import *
from setflowequation import *
from EnumDefinitions import *
from solve import *
from MatlabFuncs import *
from generic import generic

md=triangle(model(),'../Exp/Square.exp',600000)    #180000
md=setmask(md,'all','')

# Use of ispdd and isdelta18o methods
md.surfaceforcings = SMBpdd();
md.surfaceforcings.isdelta18o=1

md=parameterize(md,'../Par/SquareShelf.py')

# Add temperature, precipitation and delta18o needed to measure the surface mass balance
# creating delta18o
delta18o=numpy.loadtxt('../Data/delta18o.data')
md.surfaceforcings.delta18o=delta18o
# creating delta18oSurface
md.surfaceforcings.delta18o_surface = numpy.zeros((2,numpy.size(delta18o,axis=1)))
md.surfaceforcings.delta18o_surface[1,:] = delta18o[1,:]

# creating Present day and lgm temperatures
# Same temperature over the all region:
tmonth=numpy.ones(12)*(238.15+20.)
md.surfaceforcings.temperatures_presentday=numpy.zeros((md.mesh.numberofvertices+1,12))
md.surfaceforcings.temperatures_lgm=numpy.zeros((md.mesh.numberofvertices+1,12))
for imonth in xrange(0,12):
    md.surfaceforcings.temperatures_presentday[0:md.mesh.numberofvertices,imonth]=tmonth[imonth]
    md.surfaceforcings.temperatures_lgm[0:md.mesh.numberofvertices,imonth]=tmonth[imonth]-20.
    # Time for the last line:
    md.surfaceforcings.temperatures_presentday[md.mesh.numberofvertices,imonth]=((float(imonth)+1.)/12.)
    md.surfaceforcings.temperatures_lgm[md.mesh.numberofvertices,imonth]=((float(imonth)+1.)/12.)

# creating initialization and spc temperatures initialization and spc
md.thermal.spctemperature=numpy.mean(md.surfaceforcings.temperatures_lgm[0:md.mesh.numberofvertices,:],axis=1).reshape(-1,1)    #-10*ones(md.mesh.numberofvertices,1)
md.thermal.spctemperature=numpy.tile(md.thermal.spctemperature,(1,md.timestepping.final_time/md.timestepping.time_step))
itemp=numpy.arange(0,md.timestepping.final_time,md.timestepping.time_step)
md.thermal.spctemperature=numpy.vstack((md.thermal.spctemperature,itemp.reshape(1,-1)))

md.initialization.temperature=md.surfaceforcings.temperatures_lgm[0:md.mesh.numberofvertices,0].reshape(-1,1)   #*ones(md.mesh.numberofvertices,1)

# creating precipitation
md.surfaceforcings.precipitations_presentday=numpy.zeros((md.mesh.numberofvertices+1,12))
for imonth in xrange(0,12):
    md.surfaceforcings.precipitations_presentday[0:md.mesh.numberofvertices,imonth]=-0.4*10**(-6)*md.mesh.y+0.5
    md.surfaceforcings.precipitations_presentday[md.mesh.numberofvertices,imonth]=((float(imonth)+1.)/12.)

# time steps and resolution
md.timestepping.time_step=20.
md.settings.output_frequency=1
md.timestepping.final_time=60.

#
md.transient.requested_outputs=['default','SurfaceforcingsMonthlytemperatures']
md.extrude(3,1.)
md=setflowequation(md,'SSA','all')
md.cluster=generic('name',oshostname(),'np',1)
md=solve(md,TransientSolutionEnum())

#Fields and tolerances to track changes
field_names     =['Vx1','Vy1','Vz1','Vel1','Pressure1','Bed1','Surface1','Thickness1','Temperature1','BasalforcingsGroundediceMeltingRate1', \
						'Vx2','Vy2','Vz2','Vel2','Pressure2','Bed2','Surface2','Thickness2','Temperature2','BasalforcingsGroundediceMeltingRate2', \
						'Vx3','Vy3','Vz3','Vel3','Pressure3','Bed3','Surface3','Thickness3','Temperature3','BasalforcingsGroundediceMeltingRate3']
field_tolerances=[\
	1e-13,1e-13,1e-13,1e-13,1e-13,1e-8,1e-8,1e-13,1e-13,1e-8,\
	1e-13,1e-13,1e-13,1e-13,1e-13,1e-8,1e-8,1e-13,1e-13,1e-8,\
	1e-13,1e-13,1e-08,1e-13,1e-13,1e-8,1e-8,1e-10,1e-13,1e-8]
field_values=[\
	md.results.TransientSolution[0].Vx,\
	md.results.TransientSolution[0].Vy,\
	md.results.TransientSolution[0].Vz,\
	md.results.TransientSolution[0].Vel,\
	md.results.TransientSolution[0].Pressure,\
	md.results.TransientSolution[0].Base,\
	md.results.TransientSolution[0].Surface,\
	md.results.TransientSolution[0].Thickness,\
	md.results.TransientSolution[0].Temperature,\
	md.results.TransientSolution[0].BasalforcingsGroundediceMeltingRate,\
	md.results.TransientSolution[0].SurfaceforcingsMonthlytemperatures,\
	md.results.TransientSolution[0].SurfaceforcingsMassBalance,\
	md.results.TransientSolution[1].Vx,\
	md.results.TransientSolution[1].Vy,\
	md.results.TransientSolution[1].Vz,\
	md.results.TransientSolution[1].Vel,\
	md.results.TransientSolution[1].Pressure,\
	md.results.TransientSolution[1].Base,\
	md.results.TransientSolution[1].Surface,\
	md.results.TransientSolution[1].Thickness,\
	md.results.TransientSolution[1].Temperature,\
	md.results.TransientSolution[1].BasalforcingsGroundediceMeltingRate,\
	md.results.TransientSolution[1].SurfaceforcingsMonthlytemperatures,\
	md.results.TransientSolution[1].SurfaceforcingsMassBalance,\
	md.results.TransientSolution[2].Vx,\
	md.results.TransientSolution[2].Vy,\
	md.results.TransientSolution[2].Vz,\
	md.results.TransientSolution[2].Vel,\
	md.results.TransientSolution[2].Pressure,\
	md.results.TransientSolution[2].Base,\
	md.results.TransientSolution[2].Surface,\
	md.results.TransientSolution[2].Thickness,\
	md.results.TransientSolution[2].Temperature,\
	md.results.TransientSolution[2].BasalforcingsGroundediceMeltingRate,\
	md.results.TransientSolution[2].SurfaceforcingsMonthlytemperatures,\
	md.results.TransientSolution[2].SurfaceforcingsMassBalance,\
	]
