""" GA2STS.py - Utility to read a GA/PTHA netcdf tsunami dataset and convert
it to an Anuga sts boundary condition (netcdf) file. The supplied
GA/PTHA netcdf file contains data for stage and momentum at a series of
boundary gauge sites but also includes data for the source slide as well.
In addition the gauges are defined at lon,lats requiring conversion
to Easting and Northings before they can be used in an STS boundary file.
This script uses the scipy.io.netcdf library supplied with intel python 3.6.
It may need syntax changes to work with the unidata netcdf4 download library.
It has been created by a  netcdf user with limited python experience so is
likely not very clever/pythonic! Please feel free to improve.
It uses the utm library to convert GA lons/lats to STS easting/Northings.
This library is available at; https://pypi.org/project/utm/

V 0001 Created by Ted Rigby for Mineral Resources Tasmania August 2019
"""
import numpy as np
from scipy.io import netcdf
# note utm  is not a default module and will need to be downloaded/installed

import sys
sys.path.append("C:\\Users\\ckain\\Appdata\\Local\\Programs\\Python\\Python37-32\\Lib\\site-packages")
import utm  as utm

verbose =  True

#if verbose :
#    print ('netcdf commands available are :')
#    print (dir(netcdf))  # record available commands to assist coding
#    print ('UTM commands are             :' )
#    print (dir(utm))

#########################################################
#                   USER INPUTS                         #
#########################################################

# assign input and output netcdf file names
GAfilename  = 'newhebrides_12551.nc'
STSfilename = 'newhebrides_12551.sts'

# read the GA stations in south to north sequence to be included in the STS
# (Not not all of GA stations necessarilly included )
STSstationslist = [4,2,3,5,14,13,17,12,11,10,16,9,8,15,7,6]
NSTSstations = len(STSstationslist)

# Chech timestep to print out on screen
checktimestep = 200

#########################################################
#          READ IN THE GA nc DATA                       #
#########################################################

# open the provided GA netcdf file read only
if verbose :
    print ('Opening GA input netcdf file named ',GAfilename)
GAfilefid = netcdf.netcdf_file(GAfilename, mode='r')


# extract the data of interest
GAstations = GAfilefid.variables['station'] # (station) PTHA index 1-Nstations
GAgaugeID  = GAfilefid.variables['gaugeID'] # (station) PTHA gauge ID
GAtimes = GAfilefid.variables['time']       # (time)    PTHA times in scs from t=0
GAlons  = GAfilefid.variables['lon']        # (station) PTHA lats of station
GAlats  = GAfilefid.variables['lat']        # (station) PTHA lons of station
GAelevs = GAfilefid.variables['elev']       # (station) PTHA bath relative to MSL m
GAstages= GAfilefid.variables['stage']      # (station,event,time) PTHA m rel to msl
GAevd   = GAfilefid.variables['uh']         # (station,event,time) PTHA vxd x
GAnvd   = GAfilefid.variables['vh']         # (station,event,time) PTHA vxd y

# save the key dimensions
Nstations, Nevents, Ntimes = GAstages.shape


# print out the relevant array shapes and units for the record
if verbose :
    print ('GA variables of interest - shapes and unitswhere applicable are :')
    print ('stations ',GAstations.shape)
    print ('gaugeID  ',GAgaugeID.shape)
    print ('time     ',GAtimes.shape, GAtimes.units)
    print ('lon      ',GAlons.shape, GAlons.units)
    print ('lat      ',GAlats.shape, GAlats.units)
    print ('elevs    ',GAelevs.shape, GAelevs.units)
    print ('stages   ',GAstages.shape, GAstages.units)
    print ('evd      ',GAevd.shape, GAevd.units)
    print ('nvd      ',GAnvd.shape, GAnvd.units)

    # display station details of what we have raed in for user specified  timestep
    print('Station GaugeID   Easting      Northing     Elevation    Stage        Xmom        YMom' )
    for station in range(Nstations-1) :
        print ('%-8i %-8i %-12.3f %-12.3f %-12.3f %-12.3f %-12.3f %-12.3f' % (station, GAgaugeID[station], GAlons[station], GAlats[station], GAelevs[station], GAstages[station,0,checktimestep], GAevd[station,0,checktimestep], GAnvd[station,0,checktimestep] ))

    # dump the spatial array data for each gauge for the record
    print('GAstation GAgaugeID GAlons GAlats GAelevs')
    for station in range(0,Nstations) :
        print (GAstations[station], GAgaugeID[station], GAlons[station], GAlats[station], GAelevs[station])

#######################################################
#        CREATE A NEW FORMATTED STS FILE FROM         #
#         THE DATA READ IN FROM THE GA FILE           #
#######################################################

# open a new netcdf (sts format) file for writing
if verbose :
    print ('Opening the new (sts) netcdf file named ',STSfilename)
stsfilefid = netcdf.netcdf_file(STSfilename, mode='w')

# Assign dimensions for variables in sts file
number_of_points = stsfilefid.createDimension('number_of_points',NSTSstations )
number_of_timesteps = stsfilefid.createDimension('number_of_timesteps',Ntimes )
numbers_in_range = stsfilefid.createDimension('numbers_in_range',2)


# Create variables for storing in sts file
permutation = stsfilefid.createVariable('permutation', 'int', ('number_of_points',) )
# Note x and y are offsets relative to Lower left origin of point dataset
x = stsfilefid.createVariable('x', 'double', ('number_of_points',) )
y = stsfilefid.createVariable('y', 'double', ('number_of_points',) )
elevation = stsfilefid.createVariable('elevation', 'double', ('number_of_points',) )
elevation_range  = stsfilefid.createVariable('elevation_range', 'double', ('numbers_in_range',))
time  = stsfilefid.createVariable('time', 'double', ('number_of_timesteps',))
stage = stsfilefid.createVariable('stage', 'double', ('number_of_timesteps', 'number_of_points'))
stage_range  = stsfilefid.createVariable('stage_range', 'double', ('numbers_in_range',))
xmomentum = stsfilefid.createVariable('xmomentum', 'double', ('number_of_timesteps','number_of_points'))
xmomentum_range = stsfilefid.createVariable('xmomentum_range', 'double', ('numbers_in_range',))
ymomentum = stsfilefid.createVariable('ymomentum', 'double', ('number_of_timesteps','number_of_points'))
ymomentum_range = stsfilefid.createVariable('ymomentum_range', 'double', ('numbers_in_range',))

# assign global attributes
stsfilefid.institution = "Mineral Resources Tasmania"
stsfilefid.description = "STS Formatted Netcdf file converted from GA Tsunami Netcdf file"
stsfilefid.revision_number = "None"
stsfilefid.starttime = 0.0
stsfilefid.zone = 55
stsfilefid.false_easting  = 500000
stsfilefid.false_northing = 10000000
stsfilefid.datum = "wgs84"
stsfilefid.projection = "UTM"
stsfilefid.units = "m"

# Assign variable attributes (Hobart.sts does not seem to have assigned variable attributes

# Populate variables in sts from the GA data in new order and only include those in STS bdry
time[:] = GAtimes[:]   # not station specific

STSindex =0
for GAstationNo in STSstationslist : # populated the STS station related variables in list order(counter clockwise on bdry)
    if verbose :
        print ('assigning data for STSindex, GAstationNo, int(GAgaugeID[GAstationNo-1]) ', STSindex, GAstationNo, int(GAgaugeID[GAstationNo-1]))
    permutation[STSindex] = int(GAgaugeID[GAstationNo-1]) # (GAgaugeID ref in STS ordered list
    lat = GAlats[GAstationNo-1]
    lon = GAlons[GAstationNo-1]
    easting, northing, zone, zoneletter = utm.from_latlon(lat,lon)
    x[STSindex]           = easting
    y[STSindex]           = northing
    elevation[STSindex]   = GAelevs[GAstationNo-1]

    # now need to build STS stage xmom ymom with rearranged dimensions
    for time in range(Ntimes-1) :
        stage[time,STSindex]     = GAstages[GAstationNo-1,0,time]  # (Ntimes, Nstations)
        xmomentum[time,STSindex] = GAevd[GAstationNo-1,0,time]     # (Ntimes, Nstations)
        ymomentum[time,STSindex] = GAnvd[GAstationNo-1,0,time]     # (Ntimes, Nstations)
    STSindex = STSindex +1

# add in the derived globals
xll = min(x[:])
yll = min(y[:])
stsfilefid.xllcorner = xll
stsfilefid.yllcorner = yll

# Npw reset x[] and y[] to relative offsets for inclusion in sts
x[:] = x[:] - xll
y[:] = y[:] - yll

# print the STS dimensions out
if verbose :
    print (' STS DIMENSIONS:' )
    print (' STS number_of_points dimension    = ', stsfilefid.dimensions['number_of_points'])
    print (' STS number_of_timesteps dimension = ', stsfilefid.dimensions['number_of_timesteps'])
    print (' STS numbers_in_range dimension    = ', stsfilefid.dimensions['numbers_in_range'])

    # print the STS global attributes out
    print (' STS GLOBAL ATTRIBUTES:')
    print (' institution ',stsfilefid.institution )
    print (' description ',stsfilefid.description )
    print (' revision    ',stsfilefid.revision_number)
    print (' starttime   ',stsfilefid.starttime )
    print (' xllcorner   ',stsfilefid.xllcorner )
    print (' yllcorner   ',stsfilefid.yllcorner )
    print (' zone        ',stsfilefid.zone )
    print (' false east  ',stsfilefid.false_easting )
    print (' false north ',stsfilefid.false_northing )
    print (' datum       ',stsfilefid.datum )
    print (' projection  ',stsfilefid.projection )
    print (' units       ',stsfilefid.units )

    # display station details of what we have created for user specified  timestep
    print ('Note Values displayed for elev-xmom-ymom are for one timestep only set at ',checktimestep, ' secs')
    print ('Station Permutation Easting      Northing     Elevation   Stage        Xmom        YMom' )
    for station in range(NSTSstations-1) :
        print ('%-8i %-8i %-12.3f %-12.3f %-12.3f %-12.3f %-12.3f %-12.3f' % (station, permutation[station], x[station], y[station], elevation[station], stage[checktimestep,station], xmomentum[checktimestep,station], ymomentum[checktimestep,station] ))


