# ############################################################################################
# Program name  : calc_ref_state.py
# written by    : Gabriel Wolf, g.a.wolf@reading.ac.uk
#                 adapted from compute_reference_state_woa2009_for_Juan.m of Remi Tailleux
# last modified : 13.09.2018
# ############################################################################################

# Load necessary modules #####################################################################
import numpy as np      # 
import time             # to measure elapsed time
# read data
from get_data import get_data_woa2009 as read_data
# plotting modules / functions
from myplot import myplot_2dmap # plot of lat-lon fields
from myplot_inputs import myplot_create_pn # create plotname
# ############################################################################################

# Input 0 ####################################################################################
# Manual defined inputs
# input         : None 
# output        : dir_plot, variables
#     dir_plot  : directory, where plots will be saved
#     variables : variables necessary for calculations 
# MANUAL DEFINED OUTPUTS MISSING
# ############################################################################################
print 'MANUAL DEFINED OUTPUTS MISSING'
dim       = ('lon','lat','z','time',)
dir_plot  = '/glusterfs/inspect/users/xg911182/Code/Python/plots_test/'
variables = {'s':{},'temp':{},'grd':[],'valid':[]}
lonlat_reg = [25.0, 25.0+360.0, -90.0, 90.0] # plotting range for 2dmap
# Input 1 ####################################################################################
# Physical constants
# input         : None 
# output        : r_earth, rho0, grav, gbuo
#     r_earth   : earth radis           [m]
#     rho0      : reference density     [kg m**-3]
#     grav      : gravity               [m s**-2]
#     gbuo      : 
# PHYSICAL CONSTATS MISSING
# ############################################################################################
# constants stored in myconstants.py

# step a #####################################################################################
# Read necessary data to calculate reference state
# input         : None 
# output        : salt, theta, lon, lat, zlev, valid
#     salt      : salinity              [psu]
#     theta     : potential temperature [degC]
#     lon       : longitude             [degrees]
#     lat       : latitude              [degrees]
#     zlev      : depth of ocean        [m] ; values > 0
#     valid     : mask for gridpoints with ocean data
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print ' Read data'
grd, data = read_data(variables,dim)
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to read data:         '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step b #####################################################################################
# Define grid for data from step a
# input         : lon, lat, zlev (all from step a) 
# output        : plev
#     plev      : MISSING DEFINITION
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to define grid of data'
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to define grid:       '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step c #####################################################################################
# Calculation of density
# input         : salt, theta, plev, valid
#     salt      : output of step a
#     theta     : output of step a
#     p_lev     : output of step b
# output        : dens
#     dens      : density of water      [kg m**-3]
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to calculate density'
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to calculate density: '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step d #####################################################################################
# Calculate z-dependent ocean area
# input         : lat, r_earth,nlon, nlat, nz
#     lat       : output of step a
#     r_earth   : defined in Input 1
#     nlon, nlat, nz : MISSING - PUT THIS INTO GRD?
# output        : area_xyz
#     area_xyz  : z-dep ocean area      [m**2]
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to calculate ocean area'
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to calc. ocean area:  '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step e #####################################################################################
# Calculate reference state
# input         : dens, zlev, area_xyz
#     dens      : output of step c
#     zlev      : output of step a
#     area_xyz  : output of step d
# otput         : pp_rhor
#     pp_rhor   : interpolant of reference density
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to calculate reference state'
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to calculate reference state: '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step f #####################################################################################
# plot reference state
# input         : MISSING
# output        : None
# saved/created : plots saved in dir_plot
#     dir_plot  : output of Input 0 
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to plot reference state'
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to plot and save reference state: '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'

# step g #####################################################################################
# plot of full 2d fields
# input         : None
# output        : None
# saved/created : plots saved in dir_plot
#     dir_plot  : output of Input 0 
# ############################################################################################
beg_time = time.time()
print ''
print ' ----------------------------------------------'
print 'Insert code to plot full 2d fields'
# define plotname
list_allkeys = data.keys()
i_dep = 0
i_t   = 0
date_str = '2009'
for i_key in range(0,len(list_allkeys)):
    print 'put this plotname stuff into a function'
    # create plotname
    plotname = myplot_create_pn({'dir_plot':dir_plot,'beg_str':'Map2d',\
        'varname':list_allkeys[i_key],'zlev':[grd.z[i_dep],grd.Uz],\
        'lonlat':lonlat_reg,'date':date_str,'end_str':'.png'})
    # define plot input
    title_in = 'Depth='+'{:.1f}'.format(grd.z[i_dep])+grd.Uz
    data_in = data[list_allkeys[i_key]]['val'][:,:,i_dep,i_t]
    data_in[data_in==data[list_allkeys[i_key]]['fill_value']] = np.nan
    myplot_2dmap(data_in,grd,lonlat_range=lonlat_reg,saveplot=1,\
        title_c=title_in,plotname=plotname,unit_data=data[list_allkeys[i_key]]['units'])
    del data_in, title_in, plotname
print ' ----------------------------------------------'
end_time = time.time()
print 'Elapsed time to plot 2d fields (lat-lon): '+ '{:.2f}'.format(end_time-beg_time)+'s'
print ' ----------------------------------------------'