# 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 pylab import *
import sys
from py_rmatools import rma       
import pandas as pd             

#==========================================================#
#Input parameters
#==========================================================#
#print "Enter the RMA filename (without extension):"
#f=raw_input()

#set beginning and end years and corresponding scenario code
run='Hwq003'
startyear=1999
endyear=1999
year=range(startyear, endyear+1)

print('extracting WQ data for' + run)
    
#set run directory where RMA results are located
run_directory = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/02_Simulations/' + run +'/'
#set directory path for output files
output_directory = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/03_Results/Output_gridded1/'+ run + '/'

#set model element nodes from chainage file
#node=[666, 59,527]
#read csv file with nodes to extract data from
nodes_csv = 'H:/WRL_Projects/Hunter_CC_Modeling/Module_6/03_Results/Chainages/Hunter_nodes_gridded.csv' 
node = pd.read_csv(nodes_csv)['Hunter'].values
#==========================================================#

#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=[]
files=[]

for ii in node:
    files.append(output_directory + run + '_%d_WQ.txt' %ii)

I_const=[1]

for kk in list(enumerate(node)):
    target = open(files[kk[0]], 'w')
    target.write("Year Hour")
    for ii in I_const:
        target.write(" %d" %(ii))
    target.write("\n")
    target.close()


for jj in year:
    print(jj)
    f=run_directory + run +'_%d_SAL' %jj
    R=rma()
    R.open(f)
    while R.next():
        time.append(R.time)
        for kk in list(enumerate(node)):
            target = open(files[kk[0]], 'a')
            target.write("%i %f" %(jj,time[-1]))
            for ii in I_const:
                target.write(" %f" %(R.constit[ii][kk[1]]))
            target.write("\n")
        target.close()

###########################




#I_const=[1,2]
#filename1= output_directory + run + '_SAL.txt'
#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 +  run + '_%d' %jj
#    R=rma()
#    print(f1)
#    R.open(f1)
#    print (jj)
#    while R.next():
#        time.append(R.time)
#        for kk in list(enumerate(node)):
#            target.write("%i %f" %(jj,time[-1]))
#            target.write(" %f" %(R.constit[I_const[ii]][kk[1]]))
#        target.write('\n')
#target.close()
#
#
#
#for kk in node:
#
#    filename1='DataCR046_%d_.txt' %kk
    

#