# coding: utf-8

# In[64]:

import struct
import matplotlib.pyplot as plt
import math
from py_rmatools import rma
import re
from datetime import datetime, timedelta
import pandas as pd
import numpy as np
plt.rcParams.update({'figure.max_open_warning': 0})


# In[65]:

meshFilename = 'bub005c.rm1'
channelWidth = 100
RMAfilenames = ['BUB028','BUB029','BUB030','BUB031','BUB032']

#If RMA11
constNum = [1]

#If RMA11 3D

#Options = [4,8,1,'Max']  #Option available 'Max' , an integer correspond to the layer number to extract the results, if 2D use options = [1] - 'Percentile'
Options = ['Percentile']
percentiles = [0.5,0.8,0.9,0.95,1]
NodeLayersFile = 'NodeLayers.txt' #the node layer file was created by getting the nodes number from the RMA Outfile (just run RMA11 for a few steps)
Nnodes = 5769 #Number of nodes (including mid side node) - thisn umber also correspond to the last node number in the nodelayers file
# 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(option=1):
    f = open('{}_{}.slf'.format(run,option), 'wb')

    f.write(struct.pack(">l",80))
    strtemp='{0: >80}'.format('SERAFIN ')
    f.write(strtemp.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 idx,c in enumerate(constName):
        f.write(struct.pack(">l",32))
        strtemp='{0: <32}'.format(str(idx)+c)
        f.write(strtemp.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",startDate.year))
    f.write(struct.pack(">l",startDate.month))
    f.write(struct.pack(">l",startDate.day))
    f.write(struct.pack(">l",startDate.hour))
    f.write(struct.pack(">l",startDate.minute))
    f.write(struct.pack(">l",startDate.second))
    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):
        currentDate = datetime(R.year,1,1) + timedelta(hours = R.time)
        timeCurrentStep = currentDate - startDate
        
        f.write(struct.pack(">l",4))
        f.write(struct.pack(">f",timeCurrentStep.total_seconds()))
        f.write(struct.pack(">l",4))                        
        
        
        
        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            writeConst('X-VEL',key,f)
                
        f.write(struct.pack(">l",len(NodeList)*4))

        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            writeConst('Y-VEL',key,f)
        f.write(struct.pack(">l",len(NodeList)*4))
        
        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            writeConst('DEPTH',key,f)
        f.write(struct.pack(">l",len(NodeList)*4))
        
        f.write(struct.pack(">l",len(NodeList)*4))
        for key, value in nodeDict.items():
            writeConst('FREE SURFACE',key,f)
        f.write(struct.pack(">l",len(NodeList)*4))

        for c in constNum:
            f.write(struct.pack(">l",len(NodeList)*4))
            for key, value in nodeDict.items():
                writeConst(c,key,f)                               
            f.write(struct.pack(">l",len(NodeList)*4))
    f.close()   

def writeConst(param,key,f):
#get the surface node number
# ['X-VEL','Y-VEL','DEPTH','FREE SURFACE','Constituent 1'......] 
    
    if param == 'X-VEL':
        tempR = R.xvel
    elif param == 'Y-VEL':
        tempR = R.yvel
    elif param == 'DEPTH':
        tempR = R.depth
    elif param == 'FREE SURFACE':
        tempR = R.elevation
    else:
        tempR = R.constit[c]
 
    if key in node2nodevalue.keys():
         key = node2nodevalue[key]
    #based on the surface node number and option selected calculate the value                
    if option == 'Max':
        tempArr = []
        for n in nodelayer[key]:
            tempArr.append(tempR[n])
        tempVal1 = max(tempArr)                    
    elif isinstance(option, int):
       temppos = (option - 1 ) * 2
       temppos2 = min([temppos,len(nodelayer[key])-1])
       tempVal1 = tempR[nodelayer[key][temppos2]]
    else:
        tempVal1 = tempR[key]                                   
    #save the value in the selaphin file                
    f.write(struct.pack(">f",tempVal1)) 

def writePCTL(c,pctl):

    f = open('{}_{}_{}.slf'.format(run,c,option), 'wb')

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

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

    for idx,p in enumerate(pctl):
        f.write(struct.pack(">l",32))
        strtemp='{0: <32}'.format(str(idx)+ 'pct ' + str(p))
        f.write(strtemp.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",startDate.year))
    f.write(struct.pack(">l",startDate.month))
    f.write(struct.pack(">l",startDate.day))
    f.write(struct.pack(">l",startDate.hour))
    f.write(struct.pack(">l",startDate.minute))
    f.write(struct.pack(">l",startDate.second))
    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))

    columnName = [node2nodevalue[key] if key in node2nodevalue.keys() else key for key, value in nodeDict.items()]
    
    global dfAll
    dfAll = {}
    
    for key, value in nodeDict.items():
        if key in node2nodevalue.keys():
            key = node2nodevalue[key]
            
        dfAll[key] = np.array([])
    
    cpt = 0
    
    
    currentDate = datetime(R.year,1,1) + timedelta(hours = R.time)
    timeCurrentStep = currentDate - startDate
            
    f.write(struct.pack(">l",4))
    f.write(struct.pack(">f",timeCurrentStep.total_seconds()))
    f.write(struct.pack(">l",4)) 
    
    while R.next():
        
        tempR = R.constit[c]
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                key = node2nodevalue[key]
             
            tempArr = []
            for n in nodelayer[key]:
                tempArr.append(tempR[n])
            tempVal1 = max(tempArr)
            dfAll[key] = np.append(dfAll[key],tempVal1)

        cpt = cpt + 1
        
    for p in pctl:
        f.write(struct.pack(">l",len(NodeList)*4))
    
        for key, value in nodeDict.items():
            if key in node2nodevalue.keys():
                key = node2nodevalue[key]
            f.write(struct.pack(">f",np.quantile(dfAll[key],p)))
        
        f.write(struct.pack(">l",len(NodeList)*4))
    
def readNodeLayers(fname,Nnodes):
    fn = open(fname)
    lines = fn.readlines()
    
    surfaceNode = int(lines[0])
    nodelayer[surfaceNode] = []
    for l in lines:
        l = int(l)
        if (l != surfaceNode) & (l < Nnodes + 1):
            surfaceNode = l
            nodelayer[surfaceNode] = []
        nodelayer[surfaceNode].append(l)
        
# In[67]:

#Read mesh file and extract node (except mid node) and elements - plus convert 1D element to 2D for vizualisation
nodelayer = {}
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()
        
        temp = re.findall('.....',line)
        ElementDict[int(temp[0])] = [int(temp[i]) for i in range(1,9,2) if int(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':
        
        formatFix = (10,16,20,14,10,10)
        temp = []
        for form in formatFix:
            temp.append(line[:form])
            line = line[form:]
            
        #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]:
for run in RMAfilenames:
    for option in Options:
        RMAfilename = '{}/{}_WQ'.format(run,run)
        R=rma()
        R.open(RMAfilename)
        R.next()
        startDate = datetime(R.year,1,1) + timedelta(hours = R.time)
        if R.type==b'RMA11     ':
            constName = []
            readNodeLayers(NodeLayersFile,Nnodes)
            constName = ['X-VEL','Y-VEL','DEPTH','FREE SURFACE'] 
            for c in constNum:
                constName.append(R.constit_name[c].decode("utf-8"))
                print(c)
        
            if option == 'Percentile':
                for c in constNum:
                    writePCTL(c,percentiles)
            else:
                RMA11toSerafin(option)