from struct import *
from pylab import *
import sys

# Updated 10/11/15 BMM : Added mesh.
# Updated 10/11/15 MD : Added dictionnary initialisation in mesh.

class rma:
	def __init__(self):
			self.file='not opened yet'
			# May want to do some other things here later
			
	def open(self,filename):
			self.file=open(('%s.rma' % (filename)),'rb')
			self.header=self.file.read(1000)
			self.type=self.header[0:10]
			self.title=self.header[100:172]
			self.geometry=self.header[200:300]
			self.num_nodes=int(self.header[40:50])
			self.num_elements=int(self.header[50:60])
			if self.type==b'RMA11     ':
				self.num_constits=int(self.header[60:70])
				if len(self.header[80:90].strip())==0:
					self.num_sedlayers=0
				else:
					self.num_sedlayers=int(self.header[80:90])
				self.constit_name=[]
				self.constit_name.append("NULL")
				i=1
				#print(self.num_constits)
				while i <= self.num_constits:
					#print(i, self.num_constits, i <= self.num_constits)
					self.constit_name.append(self.header[300+(i-1)*8:308+(i-1)*8])
					if self.header[300+(i-1)*8:308+(i-1)*8]=="   SSED ":
						self.constit_name.append("  BSHEAR")
						self.constit_name.append("BedThick")
						j=1
						while j<=self.num_sedlayers:
							self.constit_name.append(" L%dThick" % j)
							j=j+1
					i=i+1
					#if self.num_sedlayers > 0:
					#		self.num_constits=self.num_constits+2+self.num_sedlayers                                
					
			# Header format in RMA2
							
			# HEADER(1:10) ='RMA2      ' 
			# HEADER(11:20)=DATEC   = Date of run
			# HEADER(21:30)=TIMEC   = Time of run
			# HEADER(31:40)=ZONEC   = Time zone
			# HEADER(41:50)=NP      = Number of nodes
			# HEADER(51:60)=NE      = Number of elements
			# HEADER(101:172)=TITLE = Title line
			# HEADER(201:300)=FNAME = File name of binary geometry

	# Header format in RMA11                
	# HEADER(1:10)  - Model 
	# HEADER(11:20) - DATEC of model runtime
	# HEADER(21:30) - TIMEC of model runtime
	# HEADER(31:40) - ZONEC of model runtime
	# HEADER(41:90) - '(5I10)' NP,NE,NQAL,IGS,LGS   or   '(5I10)' NP,NE,NQAL,ISEDS,NLAYT
	# HEADER(101:172) - TITLE  
	# HEADER(301:540) - '(30A8)' CLABL
	# HEADER(201:248) - FNAMD directory of the geometry
	# HEADER(251:298) - FNAM2 filename of the geometry

	def mesh(self,filename):
		self.meshfile=open(('%s.rm1' % (filename)),'r')
		l=self.meshfile.readline()
		while l[0:5]!=' 9999' and len(l)!=5:
			l=self.meshfile.readline()
		l=self.meshfile.readline()
		self.node_e={}
		self.node_n={}
		self.node_z={}
		while l[0:10]!='      9999' and len(l)!=10:
			ll=l.split()
			self.node_e[ll[0]]=ll[1]
			self.node_n[ll[0]]=ll[2]
			self.node_z[ll[0]]=ll[3]
			l=self.meshfile.readline()


	def next(self):
			# This reads the next timestep and populates the variables.
			# Reset all of the variables
			self.time=0.0
			self.year=0
			self.xvel=[]
			self.yvel=[]
			self.zvel=[]
			self.depth=[]
			self.elevation=[]
			
			# Add in an entry to fill position 0.  Important since RMA files start numbering nodes from 1.
			self.xvel.append(-999)
			self.yvel.append(-999)
			self.zvel.append(-999)
			self.depth.append(-999)
			self.elevation.append(-999)                
			
			if self.type==b'RMA2      ':
					# WRITE(IRMAFM) TETT,NP,IYRR,((VEL(J,K),J=1,3),K=1,NP),(WSEL(J),J=1,NP),(VDOT(3,K),K=1,NP)
					t=self.file.read(12)
					if t:
							a=unpack('fii',t)
							self.time=a[0]
							np=int(a[1])
							self.year=a[2]
							if (np!=self.num_nodes):
									print ("Warning - NP (%d) on this timestep does not match header (%d)" % (np,self.num_nodes))    
							b=unpack('%df' % 5*np, self.file.read(20*np))                                
							i=1
							while i<=np:
									self.xvel.append(b[0+(i-1)*3])
									self.yvel.append(b[1+(i-1)*3])
									self.depth.append(b[2+(i-1)*3])
									self.elevation.append(b[np*3+(i-1)])
									i=i+1
			
			if self.type==b'RMA11     ':
					self.constit=[]
					c=0
					# Start at zero to leave an empty array space.  Constits numbered 1 .. num_constits
					while c<=self.num_constits:
							self.constit.append([])
							# Put a null into the 0 position so that RMA node numbers are in the same numbered array position.
							self.constit[c].append(-999)
							c=c+1
					# READ(file1,END=100) TETT1,NQAL,NP,IYRR, ((VEL(K,J),J=1,NP),K=1,3), (wd(j),j=1,np), (wsel(j),j=1,np), ((TCON1(K,J),J=1,NP),K=1,NQAL-5)
					t=self.file.read(16)
					if t:
							a=unpack('fiii',t)
							self.time=a[0]
							nqal=int(a[1])
							np=int(a[2])
							self.year=a[3]
							if ((nqal-5)!=(self.num_constits)):
									print ("Warning - NQAL-5 (%d) on this timestep does not match header (%d)" % (nqal-5,self.num_constits))
							if (np!=self.num_nodes):
									print ("Warning - NP (%d) on this timestep does not match header (%d)" % (np,self.num_nodes))       
							b=unpack('%df' % nqal*np, self.file.read(4*nqal*np))
							i=1
							while i<=np:
									
									self.xvel.append(b[0+(i-1)*3])
									self.yvel.append(b[1+(i-1)*3])
									self.zvel.append(b[2+(i-1)*3])
									self.depth.append(b[np*3+(i-1)])
									self.elevation.append(b[np*4+(i-1)])
									c=1
									while c<=self.num_constits:
											self.constit[c].append(b[np*((c-1)+5)+(i-1)])
											c=c+1
									i=i+1
					
					
			if len(self.xvel)==1:
			   # Note that this is a 1 now as we've filled the zero array position
					return False
			else:
					return True