#==========================================================#
#Load packages
#==========================================================#
# Set working direcotry where python code is located and results are to be stored
import os
os.chdir('H:/WRL_Projects/Hunter_CC_Modeling/Code/RMA_result_explorer')
from py_rmatools_v02 import rma
import pandas as pd
import geopandas as gpd
import argparse
import glob
#==========================================================#

#==========================================================#
#Input parameters
#==========================================================#
River = 'Clarence'
foldernames =['046_CLA_CAL_01' , 'basecase_2008'] #['TWD_HYD_CAL_05', 'TWD_HYD_CAL_06']  #'base_case_2005'  #SLR_1m_2005
foldername = '046_CLA_CAL_01' 
Fishnet_ID = 'FN1'
startyear=2008 #it's critical that this matches the year of the run since the code will overwrite the year from the rma file with this one
endyear=2008   #this is technically only needed when multiple years have been run (for which the code needs to be adusted below)
year=range(startyear, endyear+1)
#==========================================================#



#==========================================================#
#Set paths to filelocations - largely automated
#==========================================================#
#load and refine node vs mesh shapefile to store the results of the statistics
Fishnet_path = glob.glob('J:/Project/wrl2018064 Fisheries RAP/04_Working/05_Modelling/RMA/HEMIP/Models/' + River + '/04_Results/Chainages/*' + Fishnet_ID + '*.shp')
Node_mesh = gpd.read_file(Fishnet_path[0])
Node_mesh = Node_mesh.loc[Node_mesh['Node'].notnull()]
Node_mesh =  pd.DataFrame(Node_mesh)
Node_mesh = Node_mesh.reset_index()
#Node_mesh.index =  Node_mesh['Field1']
#Node_mesh = Node_mesh.iloc[:,-4:]
#Node_mesh.columns = ['Node', 'X', 'Y', 'geometry']
node = Node_mesh['Node'].astype(int).values

run_directory = 'J:/Project/wrl2018064 Fisheries RAP/04_Working/05_Modelling/RMA/HEMIP/Models/' + River + '/03_Simulations/RAP_SLR/' + foldername +'/'
#set directory path for output files
output_directory = 'J:/Project/wrl2018064 Fisheries RAP/04_Working/05_Modelling/RMA/HEMIP/Models/' + River + '/04_Results/Output/'+ foldername + '_' + Fishnet_ID  + '/'  
#==========================================================#

#==========================================================#
#100% automated part of the code doing the data extraction
#==========================================================#
f = os.path.basename(glob.glob(run_directory + '*'+ str(startyear) + '*.rma')[0])[:-4] #rma file name without .rma ending

if not os.path.exists(output_directory):
    os.makedirs(output_directory)
    print('-------------------------------------------')
    print("output directory folder didn't exist and was generated")
    print('-------------------------------------------')

time=[]
xvel=[]
yvel=[]
totvel=[]
elevation=[]
depth=[]
#
#==========================================================#
#exctract elevation
#==========================================================#
print('extracting RM2 results for' + f)
filename1= output_directory + f +'_elev.txt'
filename1= output_directory + foldername +'_elev'+ '_' + Fishnet_ID + '.txt'
print(filename1)
target = open(filename1, 'w')
target.write("Year Hour ")
for inode in node:
    target.write("%i " % inode)
target.write('\n')
for jj in year:
    f1=run_directory +  f #+ '_%d' %jj
    R=rma()
    print(f1)
    R.open(f1)
    print (jj)

    while R.next():
        time.append(R.time)
        target.write("%.0f %r " %(jj,R.time))
        for inode in node:
            target.write("%f " % R.elevation[inode])
        target.write('\n')

#print (" Press any ENTER to exit")
target.close()
#f=input()
#==========================================================#

#==========================================================#
#exctract depth
#==========================================================#
filename1= output_directory + f+'_depth.txt'
print(filename1)
target = open(filename1, 'w')
target.write("Year Hour ")
for inode in node:
    target.write("%i " % inode)
target.write('\n')
for jj in year:
    print(jj)
    f1=run_directory +  f + '_%d' %jj
    R=rma()
    print(f1)
    R.open(f1)
    print (jj)

    while R.next():
        time.append(R.time)
        target.write("%.0f %r " %(jj,R.time))
        for inode in node:
            target.write("%f " % R.depth[inode])
        target.write('\n')

#print (" Press any ENTER to exit")
target.close()
#f=input()
#==========================================================#

#==========================================================#
#exctract xvel
#==========================================================#
filename1= output_directory + f+'_xvel.txt'
print(filename1)
target = open(filename1, 'w')

target.write("Year Hour ")
for inode in node:
    target.write("%i " % inode)
target.write('\n')
print (filename1)
for jj in year:
    f1=run_directory +  f + '_%d' %jj
    R=rma()
    print(f1)
    R.open(f1)
    print (jj)

    while R.next():
        time.append(R.time)
        target.write("%.0f %r " %(jj,R.time))
        for inode in node:
            target.write("%f " % R.xvel[inode])
        target.write('\n')
#print (" Press any ENTER to exit")
target.close()
#f=input()
#==========================================================#

#==========================================================#
#yvel
#==========================================================#
filename1= output_directory + f+'_yvel.txt'
print(filename1)
target = open(filename1, 'w')

target.write("Year Hour ")
for inode in node:
    target.write("%i " % inode)
target.write('\n')
print (filename1)
for jj in year:
    f1=run_directory +  f + '_%d' %jj
    R=rma()
    print(f1)
    R.open(f1)
    print (jj)

    while R.next():
        time.append(R.time)
        target.write("%.0f %r " %(jj,R.time))
        for inode in node:
            target.write("%f " % R.yvel[inode])
        target.write('\n')

#print (" Press any ENTER to exit")
target.close()
#f=input()
#==========================================================#