# coding: utf-8

# # Computing water masses volume with f2py

# Tutorial for LDEO : pick if you are running on byrd (True) or not (False)

# In[1]:


byrd=True


# In[2]:


# get the World Ocean Atlas annual T/S data

if not byrd:
    # wget in disguise in the Makefile
    get_ipython().system('make getdata')
    woadir='./'
else:
    # if running on byrd, we are gonna use files stored in ~rdussin/demo_files/
    woadir='~rdussin/demo_files/'


# In[6]:


# we need to compile the fortran code into something python can use
get_ipython().system(' source deactivate pangeo; make > /dev/null; source activate pangeo')


# In[7]:


import xarray as xr
import lib_watermass
import numpy as np


# In[8]:


#---------------------------------------------------------------
# sample data set (use make getdata to download WOA annual data)
fnameT = woadir + 'woa13_5564_t00_01.nc'
fnameS = woadir + 'woa13_5564_s00_01.nc'

woa = xr.open_mfdataset([fnameT,fnameS],decode_times=False)


# In[9]:


woa.t_an


# In[10]:


#---------------------------------------------------------------
# compute metrics (dx,dy,dz) of ocean grid cells and mask
Rearth = 6378e+3
nx = woa.lon.shape[0]
ny = woa.lat.shape[0]
nz = woa.depth.shape[0]

dxflat = (2*np.pi*Rearth/360) * (woa.lon_bnds[:,1] - woa.lon_bnds[:,0])
dyflat = (2*np.pi*Rearth/360) * (woa.lat_bnds[:,1] - woa.lat_bnds[:,0])
dzcolumn = (woa.depth_bnds[:,1] - woa.depth_bnds[:,0]).values

dxflat2d, dyflat2d = np.meshgrid(dxflat,dyflat)
lon2d, lat2d = np.meshgrid(woa.lon,woa.lat)

dx = dxflat2d * np.cos(2*np.pi*lat2d/360)
dy = dyflat2d

dz = np.empty((nz,ny,nx))
for jj in np.arange(ny):
        for ji in np.arange(nx):
                dz[:,jj,ji] = dzcolumn


# In[11]:


#---------------------------------------------------------------
# compute mask from missing value
mask = np.ma.array(np.ones(woa.t_an.squeeze().shape))
mask[np.isnan(woa.t_an.values.squeeze())] = 0


# In[12]:


# Quick check on the mask
get_ipython().run_line_magic('pylab', 'inline')
import matplotlib.pylab as plt
plt.figure(figsize=[12,12])
plt.pcolormesh(mask[0,:,:]) ; plt.colorbar()
plt.show()


# In[13]:


#---------------------------------------------------------------
# sanity checks
# compute water volume between 50C and 55C (never observed)
wmass_testimpossible = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                              woa.t_an.transpose(),                                                              woa.s_an.transpose(),                                                              50.,55.,0.,40.)


# In[14]:


# compute all possible volume
wmass_testall = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),                                                       woa.s_an.transpose(),                                                       -10.,100.,0.,50.)

# and compare to volume from metrics
wmass_from_metrics = (dx * dy * mask * dz).sum()

print('---Control checks---')
print('Impossible water mass, volume =', wmass_testimpossible)
print('All possible water, volume =', wmass_testall)
print('volume from scale factors = ', wmass_from_metrics)


# In[15]:


#---------------------------------------------------------------
# now let's learn some stuff
# compute volume of water 20C < T < 40C
wmass_20to40C = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),                                                       woa.s_an.transpose(),                                                       20.,40.,0.,40.)
# compute volume of water 10C < T < 20C
wmass_10to20C = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),woa.s_an.transpose(),                                                       10.,20.,0.,40.)
# compute volume of water 0C < T < 10C
wmass_0to10C = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                      woa.t_an.transpose(),woa.s_an.transpose(),                                                      0.,10.,0.,40.)

print('---Volume of ocean in temperature ranges---')
print('The percentage of ocean waters 20C < T < 40C is ', 100 * wmass_20to40C / wmass_from_metrics, '%')
print('The percentage of ocean waters 10C < T < 20C is ', 100 * wmass_10to20C / wmass_from_metrics, '%')
print('The percentage of ocean waters  0C < T < 10C is ', 100 * wmass_0to10C / wmass_from_metrics, '%')


# In[16]:


# note on I/O and dimensions

#---------------------------------------------------------------

wmass_20to40C = lib_watermass.volume_watermass_from_ts(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),                                                       woa.s_an.transpose(),                                                       20.,40.,0.,40.)

print('Solution from v1 is ', 100 * wmass_20to40C / wmass_from_metrics, '%')

wmass_20to40C = lib_watermass.volume_watermass_from_ts_v2(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),                                                       woa.s_an.transpose(),                                                       20.,40.,0.,40.,nx,ny,nz)


print('Solution from v2 is ', 100 * wmass_20to40C / wmass_from_metrics, '%')

wmass_20to40C = lib_watermass.volume_watermass_from_ts_v2(dx.transpose(),dy.transpose(),dz.transpose(),                                                       woa.t_an.transpose(),                                                       woa.s_an.transpose(),                                                       20.,40.,0.,40.)

print('Solution from v2 without nx,ny,nz is ', 100 * wmass_20to40C / wmass_from_metrics, '%')

wmass_20to40C = lib_watermass.volume_watermass_from_ts_v3(dx,dy,dz,                                                       woa.t_an,                                                       woa.s_an,                                                       20.,40.,0.,40.)

print('Solution from v3 without nx,ny,nz is ', 100 * wmass_20to40C / wmass_from_metrics, '%')