# coding: utf-8

# In[64]:

import struct
import matplotlib.pyplot as plt
import math
from py_rmatools import rma

plt.rcParams.update({'figure.max_open_warning': 0})


# In[65]:

meshFilename = 'HWQ010.rm1'
channelWidth = 100
RMAfilename = 'HWQ277_2010_WQ'

#If RMA11
constNum = [1,2,3,5,6,7,8,10,11,13,14,15,18]


# In[66]:

def isElementOneD(nodelist):
    if len(nodelist) == 2:
        return True
    return False

def isElementSquare(nodelist):
    if len(nodelist) == 4:
        return True
    return False

def square2Triangle(ElementNum):
    nodelist = ElementDict[ElementNum]
    if isElementSquare(nodelist):
        ElementDict[ElementNum] = [nodelist[0], nodelist[1], nodelist[2]]
        ElementList.append(max(ElementList) + 1)
        ElementDict[ElementList[-1]]= [nodelist[0], nodelist[2], nodelist[3]]
        
def oneD2triangle(ElementNum):
    if isElementOneD(ElementDict[ElementNum]):
        nAe = ElementDict[ElementNum][0] #nAe Node A existing
        nBe = ElementDict[ElementNum][1]
        
        if not nAe in node1Dduplicate: node1Dduplicate[nAe] = []
        if not nBe in node1Dduplicate: node1Dduplicate[nBe] = []
        
        xA = nodeDict[nAe][0]
        xB = nodeDict[nBe][0]
        yA = nodeDict[nAe][1]
        yB = nodeDict[nBe][1]
        
        normalVec = [-(yB - yA),(xB - xA)]
        dist = math.sqrt(normalVec[0]**2 + normalVec[1]**2)
        normalVec[0] = normalVec[0] / dist
        normalVec[1] = normalVec[1] / dist
        xA2 = xA + channelWidth * normalVec[0]
        xB2 = xB + channelWidth * normalVec[0]
        yA2 = yA + channelWidth * normalVec[1]
        yB2 = yB + channelWidth * normalVec[1]
        

        nA = max(NodeList) + 1
        nB = max(NodeList) + 2
        
        node1Dduplicate[nAe].append(nA)
        node1Dduplicate[nBe].append(nB) 
        
        node2nodevalue[nA] = nAe
        node2nodevalue[nB] = nBe
        
        
        NodeList.append(nA)
        NodeList.append(nB)
        nodeDict[nA] = [xA2, yA2, -1.01]
        nodeDict[nB] = [xB2, yB2, -1.01]
        
        newEle = max(ElementList) + 1
        ElementList .append(newEle)
        ElementDict[ElementNum] = [nAe,  nA,  nBe]
        ElementDict[newEle] = [nA,  nB,  nBe]
        
def RMA11toSerafin():
    f = open('{}.slf'.format(RMAfilename), 'wb')

    f.write(struct.pack(">l",80))
    str='{0: >80}'.format('SERAFIN ')
    f.write(str.encode('ascii'))
    f.write(struct.pack(">l",80))

    f.write(struct.pack(">l",8))
    f.write(struct.pack(">l",len(constNum)))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",8))

    for c in constName:
        f.write(struct.pack(">l",32))
        str='{0: <32}'.format(c)
        f.write(str.encode('ascii'))
        f.write(struct.pack(">l",32))


    f.write(struct.pack(">l",40))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",40))

    f.write(struct.pack(">l",24))
    f.write(struct.pack(">l",R.year))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",24))

    f.write(struct.pack(">l",16))
    f.write(struct.pack(">l",len(ElementList)))
    f.write(struct.pack(">l",len(NodeList)))
    f.write(struct.pack(">l",3))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",16))


    f.write(struct.pack(">l",len(ElementList)*3*4))
    for el in ElementList:
        for nd in ElementDict[el]:
            f.write(struct.pack(">l",nodeOrdered[nd])) 
    f.write(struct.pack(">l",len(ElementList)*3*4))

    f.write(struct.pack(">l",len(NodeList)))
    for i in range(0,len(NodeList)):
        f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",len(NodeList)))

    f.write(struct.pack(">l",len(NodeList)*4))
    for key, value in nodeDict.items():
        f.write(struct.pack(">f",value[0]))
    f.write(struct.pack(">l",len(NodeList)*4))

    f.write(struct.pack(">l",len(NodeList)*4))
    for key, value in nodeDict.items():
        f.write(struct.pack(">f",value[1]))
    f.write(struct.pack(">l",len(NodeList)*4))



    while R.next():
    #for i in range(3):
        f.write(struct.pack(">l",4))
        f.write(struct.pack(">f",R.time * 3600))
        f.write(struct.pack(">l",4))

        for c in constNum:
            f.write(struct.pack(">l",len(NodeList)*4))
            for key, value in nodeDict.items():
                if key in node2nodevalue.keys():
                    f.write(struct.pack(">f",R.constit[c][node2nodevalue[key]]))
                else:
                    f.write(struct.pack(">f",R.constit[c][key]))

            f.write(struct.pack(">l",len(NodeList)*4))

        #R.next()

    f.close()   

    
    
def RMA2toSerafin():
    f = open('{}.slf'.format(RMAfilename), 'wb')

    f.write(struct.pack(">l",80))
    str='{0: >80}'.format('SERAFIN ')
    f.write(str.encode('ascii'))
    f.write(struct.pack(">l",80))

    f.write(struct.pack(">l",8))
    f.write(struct.pack(">l",len(constName)))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",8))

    for c in constName:
        f.write(struct.pack(">l",32))
        str='{0: <32}'.format(c)
        f.write(str.encode('ascii'))
        f.write(struct.pack(">l",32))


    f.write(struct.pack(">l",40))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",40))

    f.write(struct.pack(">l",24))
    f.write(struct.pack(">l",R.year))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",24))

    f.write(struct.pack(">l",16))
    f.write(struct.pack(">l",len(ElementList)))
    f.write(struct.pack(">l",len(NodeList)))
    f.write(struct.pack(">l",3))
    f.write(struct.pack(">l",1))
    f.write(struct.pack(">l",16))


    f.write(struct.pack(">l",len(ElementList)*3*4))
    for el in ElementList:
        for nd in ElementDict[el]:
            f.write(struct.pack(">l",nodeOrdered[nd]))
    f.write(struct.pack(">l",len(ElementList)*3*4))

    f.write(struct.pack(">l",len(NodeList)))
    for i in range(0,len(NodeList)):
        f.write(struct.pack(">l",0))
    f.write(struct.pack(">l",len(NodeList)))

    f.write(struct.pack(">l",len(NodeList)*4))
    for key, value in nodeDict.items():
        f.write(struct.pack(">f",value[0]))
    f.write(struct.pack(">l",len(NodeList)*4))

    f.write(struct.pack(">l",len(NodeList)*4))
    for key, value in nodeDict.items():
        f.write(struct.pack(">f",value[1]))
    f.write(struct.pack(">l",len(NodeList)*4))



    while R.next():

        f.write(struct.pack(">l",4))
        f.write(struct.pack(">f",R.time * 3600))
        f.write(struct.pack(">l",4))


        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                f.write(struct.pack(">f",R.xvel[node2nodevalue[key]]))
            else:
                f.write(struct.pack(">f",R.xvel[key]))
        f.write(struct.pack(">l",len(NodeList)*4))

        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                f.write(struct.pack(">f",R.yvel[node2nodevalue[key]]))
            else:
                f.write(struct.pack(">f",R.yvel[key]))
        f.write(struct.pack(">l",len(NodeList)*4))
        
        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                f.write(struct.pack(">f",R.depth[node2nodevalue[key]]))
            else:
                f.write(struct.pack(">f",R.depth[key]))
        f.write(struct.pack(">l",len(NodeList)*4))
        
        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                f.write(struct.pack(">f",R.elevation[node2nodevalue[key]]))
            else:
                f.write(struct.pack(">f",R.elevation[key]))
        f.write(struct.pack(">l",len(NodeList)*4))
        
        
        R.next()

    f.close()   


# In[67]:

#Read mesh file and extract node (except mid node) and elements - plus convert 1D element to 2D for vizualisation
NodeList = [] 
ElementList = []
ElementDict = {}
nodeDict = {}
node1Dduplicate = {}    #Original Number: List of Duplicates
node2nodevalue = {}     #link between the node number and the node value to use 
                        #(e.g. if node 10 is a 1D node: 10 is not duplicate so {1:1}, 
                        #but node 2050 (duplicate of 10) (1D to 2D) the value of the duplicated 
                        #node will be the same as the original so we might have {2050: 10})
            
with open(meshFilename) as f:
    line = f.readline()
    line = f.readline()
    line = f.readline()
    line = f.readline()
    
    cpt = 1
    while line and line != ' 9999\n':
        temp = line.split()
        ElementDict[int(temp[0])] = [int(temp[i]) for i in range(1,9,2) if temp[i] != '0' and int(temp[9]) < 100]
        ElementList.append(int(temp[0]))
        line = f.readline()
          
    for key, value in ElementDict.items():
        NodeList.extend(value)
    
    NodeList = list(set(NodeList))

    line = f.readline()
    while line and line != '      9999\n':
        temp = line.split()
        if int(temp[0]) in NodeList:
            nodeDict[int(temp[0])] = [float(temp[1]),float(temp[2]),float(temp[3])]
        line = f.readline()


for e in ElementList:
    oneD2triangle(e)
    square2Triangle(e)
    
for key in list(ElementDict): #Remove Special Element 902.....
    if len(ElementDict[key]) != 3:
        print(key, ElementDict[key])
        ElementDict.pop(key)
        ElementList.remove(key)
        
nodeOrdered = {}
cpt = 1
for key, value in nodeDict.items():
    nodeOrdered[key] = cpt
    cpt +=1



# # Open and Read First Step of the RMA File and Save a Serafin

# In[72]:

R=rma()
R.open(RMAfilename)
R.next()

if R.type==b'RMA11     ':
    constName = []
    for c in constNum:
        constName.append(R.constit_name[c].decode("utf-8"))
    RMA11toSerafin()
        
if R.type==b'RMA2      ':
    constName = ['X-VEL','Y-VEL','DEPTH','FREE SURFACE']  
    RMA2toSerafin()