# ############################################################################################
# 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_xyz # calc horizontal area
from gw_ocean_refstate import refstate_meandensity, refstate_sorted_stheta # calc. ref. states
# 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
import pickle # to save reference state and create anomalies
# Further modules
import os # to execute terminal commands
# ############################################################################################

# Input 0 ####################################################################################
# Manual defined inputs
# input         : None 
# output        : dim, dir_plot, saveplot, variables, REST MISSING
#     dim       : prefered dimensions for matrices     [list]
#     dir_plot  : directory, where plots will be saved [string]
#     saveplot  : defines if plots are being saved     [0 or 1]
#     variables : variables necessary for calculations [MISSING]
# MANUAL DEFINED OUTPUTS MISSING
# ############################################################################################

print '----------------------------------------------'
print 'Read manual defined input data'
import input_calc_ref_state
# information about input data
print '  *** MANUAL DEFINED OUTPUTS MISSING ***'
dim        = ('lon','lat','z','time',)
variables  = {'s':{},'theta':{},'grd':[],'valid':[]}
print '  Save plots? Enter 1 for YES or 0 for NO.'
saveplot   = input('  -> USER INPUT: ')
print '  Define case with string (used to specify plots and data), e.g. \'_id000\''
case_id    = input('  -> USER INPUT: ')
print '----------------------------------------------'

# 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)
if grd.lon[-1]+grd.lon[0] >= 1:
   grd.symlon = 1
else:
   grd.symlon = 0
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']           = np.zeros_like(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)],reference=ref_EOS)
# Delete 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_xyz(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         : area_xyz, dens, _t, zlev, ztarget
#     dens      : output of step c
#     i_t       : MISSING - DONE LOCALLY HERE
#     zlev      : output of step a
#     area_xyz  : output of step d
#     ztarget   : target depth          [m], 1d - DONE LOCALLY HERE
# otput         : pp_rhor, pp_rhor_botup, pp_rhor_topdn
#     pp_rhor   : interpolant for ref. density (hor. av. rho)
#     pp_rhor_botup : interpolant for ref. density (sorted in th-s-space, calc. bottom-up)
#     pp_rhor_topdn : interpolant for ref. density (sorted in th-s-space, calc. top-down)
# ############################################################################################

# THIS MUST BE DONE AUTOMATICALLY
i_t       = 0

beg_time  = time.time()
print ''
print '----------------------------------------------'
print 'Calculate reference state for mean density'

# Define input data
mask      = data['rho']['valid'][:,:,:,i_t]

# define target depth
n_target      = grd.Nz+9
ztarget       = np.zeros(n_target)
ztarget[0:10] = np.arange(0,10)
ztarget[10:]  = grd.z[1:]

# calculation of reference profiles
pp_rhor   = refstate_meandensity(data['rho']['val'][:,:,:,i_t]*mask,grd,area_xyz*mask)
end_time  = time.time()
beg_time2 = time.time()

# Define input data
rho  = data['rho']['val']
s    = data['s']['val']
th   = data['theta']['val']
mask = data['rho']['valid']*data['s']['valid']*data['theta']['valid']

# calculation of reference profiles
print 'Calculate sorted density reference state'
pp_rhor_botup, pp_rhor_topdn = refstate_sorted_stheta(th, s, rho, mask, ztarget, grd,\
    ref_EOS=ref_EOS, sbin=[34.0,35.0,0.002], thbin=[-2.6,3.0,0.005],\
    test=True, case_str=case_id)

# Save output of different reference states
print 'Save reference state'
with open('data_ref_state' + case_id + '.pickle', 'wb') as output:
    pp_refstate = [pp_rhor_mean, pp_rhor_botup, pp_rhor_topdn]
    pickle.dump(pp_refstate, output)

# Deallocation
del th, s, rho

print '----------------------------------------------'
end_time2 = time.time()
print 'Elapsed time to calculate reference state: '
print '-> mean density (hor. averaged):           '+ '{:.2f}'.format(end_time-beg_time)+'s'
print '-> sorted density in S-Theta diagram:      '+ '{:.2f}'.format(end_time2-beg_time2)+'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 'Plot z-dependent reference state'

# Create plotname
date_str = 'WOA2009'
pn_full  = myplot_create_pn({'dir_plot':dir_plot,'beg_str':'Refstate',\
    'varname':'rho','date':date_str,'end_str':case_id+'.png'})
print 'Plot deviation of ref-state to a previous calculated ref-state?'
plot_dev = input('-> USER INPUT: ')
if plot_dev:
    print 'Please enter one of the listed filenames as comparison reference state.'
    os.system("ls data_ref_state*.pickle")
    fn_refstate = input('-> USER INPUT: ')
    fn_diff     = [li for li in list(difflib.ndiff(fn_refstate,\
        'data_ref_state.pickle')) if li[0] != ' ']
    pn_anom     = myplot_create_pn({'dir_plot':dir_plot,'beg_str':'Refstate_dev'+fn_diff,\
        'varname':'rho','date':date_str,'end_str':case_id+'.png'})

# Define plot input
xxx = np.zeros([3,n_target])
xxx[0,:] = pp_rhor_mean(ztarget)
xxx[1,:] = pp_rhor_botup(ztarget)
xxx[2,:] = pp_rhor_topdn(ztarget)
yyy = np.tile(-ztarget,[3,1])
if plot_dev:
    xxx_anom = np.zeros([3,n_target])
    with open(fn_refstate, 'rb') as data_ref:
        dummy         = pickle.load(data_ref)
        xxx_anom[0,:] = xxx[0,:] - dummy[0](ztarget)
        xxx_anom[1,:] = xxx[1,:] - dummy[1](ztarget)
        xxx_anom[2,:] = xxx[2,:] - dummy[2](ztarget)

# Plotting
myplot_1d(xxx,yyy,FS=22,col_c=('b','r','k'),saveplot=saveplot,\
    label_c=('mean','bottom-up','top-down'),xlabel_c='rho_ref [kg m**-3]',\
    ylabel_c='z_ref [m]',plotname=pn_full)
if plot_dev:
    myplot_1d(xxx,yyy,FS=22,col_c=('b','r','k'),saveplot=saveplot,\
        label_c=('mean','bottom-up','top-down'),xlabel_c='rho_ref [kg m**-3]',\
        ylabel_c='z_ref [m]',plotname=pn_anom)

# Deallocation
if plot_dev:
    del plot_anom, fn_refstate, fn_diff, dummy, xxx_anom
del plot_dev, pn_full, xxx, yyy

print '----------------------------------------------'
end_time = time.time()
print 'Elapsed time to plot 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
# ############################################################################################

if plot_2d:
    beg_time = time.time()
    print '\n----------------------------------------------'
    print 'Plot 2d fields (lat-lon)'

    # Define Input for all following 2d-plots
    list_allkeys = plot_var.keys() # all variables
    i_dep    = 0
    i_t      = 0
    date_str = 'WOA2009'
    
    # Loop over all variables
    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':case_id+'.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

        # Plotting
        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'])

        # Deallocation
        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 = 'WOA2009'
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':case_id+'.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':case_id+'.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 '----------------------------------------------'