# ############################################################################################
# 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 : 19.09.2018
# ############################################################################################
# More information in the README.txt file
# Necessary user input:
#     manual user input under Input 0
#     reading of data under step a
# -> everything else should work by using this Code
# Table of content of this programme:
# - Input 0 :
# - Input 1 :
# - step a  :
# - step b  :
# - step c  :
# - step d  :
# - step e  :
# - step f  :
# - step g  :
# - step h  :
# ############################################################################################

# 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
# calculation modules
from mycalc_ocean import calc_rho_from_theta # calc of density from theta
from mycalc_ocean import calc_area_lat # calc horizontal area
# plotting modules / functions
from myplot import myplot_2dmap # plot of lat-lon fields
from myplot import myplot_2d    # plot of crossections 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
# ############################################################################################
# information about input data
print 'MANUAL DEFINED OUTPUTS MISSING'
dim        = ('lon','lat','z','time',)
variables  = {'s':{},'theta':{},'grd':[],'valid':[]}
# plotting input
plot_var   = {'rho':[]}
saveplot   = input('Do you want to save the plots? Enter 1 for yes or 0 for no')
dir_plot   = '/glusterfs/inspect/users/xg911182/Code/Python/plots_test/'
lonlat_reg = [25.0, 25.0+360.0, -90.0, 90.0] # plotting range for 2dmap
plot_2d    = 1 # plot 2d fields (lat-lon)
plot_cs    = 1 # plot crosssections
cs_plot    = {'lon':[-47.5,172.5],'lat':[0]} # crosssections
# 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      : 
# ############################################################################################
from myconstants import rho0, grav, r_earth

# 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        : lat_xy, lat_yz, lat_xyz, lon_xy, lon_xyz, p_xyz, plev, z_yz
#     lat_*     : meshgrid for latitude in x/y/z-dim  [degrees]
#     lon_*     : meshgrid for longitude in x/y/z-dim [degrees]
#     p_xyz     : meshgrid for pressure in xyz-dim    [dbar]
#     plev      : gauge pressure (p_abs-10.1325 dbar) [dbar]
#     z_yz      : meshgrid for depth in yz-dim        [m]
# ############################################################################################
beg_time = time.time()
print ''
print '----------------------------------------------'
print 'Create 2d and 3d meshgrids'
try:
    plev  = grd.p
except:
    plev   = (rho0*grav/1e4)*grd.z
    grd.p  = plev
    grd.Up = 'dbar'
    grd.Np = len(plev)
lat_xy, lon_xy          = np.meshgrid(grd.lat,grd.lon)
p_yz, lat_yz            = np.meshgrid(grd.z, grd.lat)
lat_xyz, lon_xyz, p_xyz = np.meshgrid(grd.lat, grd.lon, grd.p)
print '----------------------------------------------'
end_time = time.time()
print 'Elapsed time to create meshgrids:  '+ '{:.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
#     plev      : output of step b
# output        : dens
#     dens      : density of water      [kg m**-3]
# ############################################################################################
beg_time = time.time()
print ''
print '----------------------------------------------'
print 'Calculate density'
# Theta mask for valid values
mask = data['theta']['valid']
# Allocate density
data['rho']                  = {}
data['rho']['val']           = data['theta']['val']
data['rho']['fill_value']    = data['theta']['fill_value']
data['rho']['standard_name'] = 'sea_water_density'
data['rho']['valid']         = mask
data['rho']['units']         = 'kg m**-3'
# calculate density
data['rho']['val'][mask] = calc_rho_from_theta(data['s']['val'][mask],\
    data['theta']['val'][mask],p_xyz[np.squeeze(mask)])
# De;ete unnecessary variables
del mask
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 'Calculate z-dependent ocean area'
area_xyz = calc_area_lat(grd,r_earth)
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         : data, date_str, dir_plot, grd, i_dep, i_t, lonlat_reg, plot_var
#     data      : output of step a
#     date_str  : MISSING - DONE LOCALLY HERE
#     dir_plot  : output of Input 0
#     grd       : output of step a
#     i_dep     : MISSING - DONE LOCALLY HERE
#     i_t       : MISSING - DONE LOCALLY HERE
#     lonlat_reg: output of Input 0
#     plot_var  : output of Input 0
# output        : None
# saved/created : plots saved in dir_plot
# ############################################################################################
beg_time = time.time()
print ''
print '----------------------------------------------'
print 'Plot 2d fields (lat-lon)'
list_allkeys = plot_var.keys() # all variables
i_dep    = 0
i_t      = 0
date_str = '2009'
for i_key in range(0,len(list_allkeys)):
    # 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 '----------------------------------------------'

# step h #####################################################################################
# plot of crossection fields
# input         : cs_plot, data, date_str, dir_plot, grd, i_t, plot_cs, plot_var, saveplot
#     cs_plot   : output of Input 0
#     data      : output of step a
#     date_str  : MISSING - DONE LOCALLY HERE
#     dir_plot  : output of Input 0
#     grd       : output of step a 
#     i_t       : MISSING - DONE LOCALLY HERE
#     plot_cs   : output of Input 0
#     plot_var  : output of Input 0
#     saveplot  : output of Input 0
# output        : None
# saved/created : plots saved in dir_plot
# ############################################################################################
date_str = '2009'
i_t      = 0
if plot_cs:
    beg_time = time.time()
    print ''
    print '----------------------------------------------'
    print 'Plot lon-p and lat-p crossections'
    list_allkeys = plot_var.keys() # all variables
    for i_key in range(0,len(list_allkeys)):
        for i_cs in range(0,len(cs_plot['lat'])):
            # load fixed latitude and associated index for grd.lat
            lat_fix  = cs_plot['lat'][i_cs]
            i_lat    = abs(grd.lat-cs_plot['lat'][i_cs]).argmin()
            # create plotname
            plotname = myplot_create_pn({'dir_plot':dir_plot,'beg_str':'CS',\
                'varname':list_allkeys[i_key],'lonlat':[0,360,lat_fix,lat_fix],\
                'date':date_str,'end_str':'.png'})
            # define plot input
            title_in = 'Depth='+'{:.1f}'.format(grd.z[i_dep])+grd.Uz
            cbarlabel_in = data[list_allkeys[i_key]]['standard_name'] + \
                ' ['+data[list_allkeys[i_key]]['units']+']'
            data_in = np.transpose(np.squeeze(data[list_allkeys[i_key]]['val'][:,i_lat,:,i_t]),[1,0])
            data_in[data_in==data[list_allkeys[i_key]]['fill_value']] = np.nan
            myplot_2d(grd.lon,-grd.z,data_in,saveplot=saveplot,FS=14,d_xtick=45,\
               xlabel_c='Longitude ['+grd.Ulon+']',ylabel_c='Depth ['+grd.Uz+']',\
               title_c=title_in,plotname=plotname,cbarlabel_c=cbarlabel_in)
            del data_in, lat_fix, i_lat, title_in, plotname, cbarlabel_in
        for i_cs in range(0,len(cs_plot['lon'])):
            # load fixed longitude and associated index for grd.lon
            lon_fix  = cs_plot['lon'][i_cs]
            i_lon    = abs(grd.lon-cs_plot['lon'][i_cs]).argmin()
            # create plotname
            plotname = myplot_create_pn({'dir_plot':dir_plot,'beg_str':'CS',\
                'varname':list_allkeys[i_key],'lonlat':[lon_fix,lon_fix,-90,90],\
                'date':date_str,'end_str':'.png'})
            # define plot input
            title_in = 'Depth='+'{:.1f}'.format(grd.z[i_dep])+grd.Uz
            cbarlabel_in = data[list_allkeys[i_key]]['standard_name'] + \
                ' ['+data[list_allkeys[i_key]]['units']+']'
            data_in = np.transpose(np.squeeze(data[list_allkeys[i_key]]['val'][i_lon,:,:,i_t]),[1,0])
            data_in[data_in==data[list_allkeys[i_key]]['fill_value']] = np.nan
            myplot_2d(grd.lat,-grd.z,data_in,saveplot=saveplot,FS=14,d_xtick=45,\
               xlabel_c='Latitude ['+grd.Ulat+']',ylabel_c='Depth ['+grd.Uz+']',\
               title_c=title_in,plotname=plotname,cbarlabel_c=cbarlabel_in)
            del data_in, lon_fix, i_lon, title_in, plotname, cbarlabel_in
    print '----------------------------------------------'
    end_time = time.time()
    print 'Elapsed time to plot crossections: '+ '{:.2f}'.format(end_time-beg_time)+'s'
    print '----------------------------------------------'