#==========================================================#
#Load packages
#==========================================================#
# Set working direcotry where python code is located and results are to be stored
import os
os.chdir('J:/Project/wrl2018064 Fisheries RAP/04_Working/05_Modelling/RMA/HEMIP/RMA_result_explorer_code')
from py_rmatools_v02 import rma
import geopandas as gpd
import pandas as pd
import glob
from pyrma import pyrma
#==========================================================#

#==========================================================#
#Input parameters
#==========================================================#
#set beginning and end years and corresponding scenario code
River = 'Tweed'
foldernames =['TWD_HYD_CAL_03','TWD_HYD_CAL_04' , 'TWD_HYD_CAL_05', 'TWD_HYD_CAL_06']  #'base_case_2005'  #SLR_1m_2005
Order_ID = 'Core_order' #choose order column name

startyear=2017 #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=2017   #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)
#==========================================================#



#==========================================================#
#After here should be mostly automated
#==========================================================#
for foldername in foldernames:
    #set run directory where RMA results are located
    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 + '_' + Order_ID  + '/'
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)
    
    #load MHL gauge shapefile to store the results of the statistics
    MHL_gauges_df = gpd.read_file('J:/Project/wrl2018064 Fisheries RAP/04_Working/05_Modelling/RMA/HEMIP/Global_Data/GIS/MHL_gauges_received_MGA56.shp')
    MHL_gauges_df = MHL_gauges_df.loc[MHL_gauges_df['River'] == River]
    MHL_gauges_df = MHL_gauges_df.loc[MHL_gauges_df[Order_ID].notnull()]
    MHL_gauges_df =  pd.DataFrame(MHL_gauges_df )
    MHL_gauges_df = MHL_gauges_df.reset_index()
    MHL_gauges_df['NrstNode'] = 0
    
    # Load RMA mesh and find nearest node for each MHL gauge as basis for the data extraction
    print('Reading RMA mesh file and finding the nearest node for the following XY locations')
    meshpath = glob.glob(run_directory + '*.rm1')
    nodes, elements = pyrma.loadMesh(meshpath[0])
    mesh = pd.DataFrame(elements, index=[0]).transpose()
    mesh.columns = ['name']
    mesh['centroid'] = [e.centroid for e in elements.values()]
    node = []
    for i in range(0,len(MHL_gauges_df),1):
        # Find nearest element in RMA mesh
        x = MHL_gauges_df['Easting'][i]
        y = MHL_gauges_df['Northing'][i]
        print(x,y)
        mesh['calc'] = [pyrma.point(x, y).dist(c) for c in mesh['centroid']]
        idx = mesh['calc'].idxmin()
        node.append(idx)
    
    MHL_gauges_df['NrstNode'] = node
    MHL_gauges_df.to_csv(output_directory + '_Gauges_with_nodes.csv')
    #==========================================================#
    
    #==========================================================#
    #100% automated part of the code doing the data extraction
    #==========================================================#
    #define rma file from which data will be extracted
    f = os.path.basename(glob.glob(run_directory + '*'+ str(startyear) + '*.rma')[0])[:-4] #rma file name without .rma ending
    
    #output_directory = 'J:/Project wrl2016032/NARCLIM_Raw_Data/Extracted'
    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 + foldername +'_elev'+ '_' + Order_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 +  foldername+'_depth'+ '_' + Order_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:
        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 + foldername+'_xvel'+ '_' + Order_ID + '.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 + foldername+'_yvel'+ '_' + Order_ID + '.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()
    #==========================================================#