#!/usr/bin/env python
# coding: utf-8

# This notebooks is used to develope and test a function that will calulate the depth average of a model variable.

# # Planning

# * Use np.trapz to integrate
# * Average by dividing by the depth
# 
# Finding the depth might be tricky. We could use the bathymetry but I think there are differences between the bathymetry file and the actual depth of the grid cell. I can think of two approaches
# 
# 1. Take advantage of masking. Look up index of masked point and evaluate the depth.
# 2. Use the batyhmetry to look up the grid point depth.
# 
# My plan is to use the first approach.

# # Loading

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import os


# # Function

# In[2]:


def depth_average(var,depths,depth_axis):
    """Average over depth using the trapezoid rule. The variable should be masked to apply this function. 
    The depth is calcluated based on masking. 
    If the variable is not masked then the maximum depth of the entire domain is used. 
    
    :arg var: variable to average
    :type var: masked numpy array
    
    :arg depths: the depths associated with var
    :type depths: numpy array
    
    :arg depth_axis: The axis in var associated with depth
    :type depth_axis: int
    
    :returns: avg, the depth averaged var.
    
    """
    # Make sure depths is an array and not an netcf variable.
    de=np.array(depths)
    # Integrate, the easy part
    integral = np.trapz(var,x=de,axis=depth_axis)
    # Find depth for averaging
    # Need to expand the depths array to same shape as the variable. This is really awkward..
    for n in np.arange(var.ndim-1):
        de=de[:,np.newaxis]
    roll=np.rollaxis(var,depth_axis)
    expanded_depths=  de +np.zeros(roll.shape)
    expanded_depths=np.rollaxis(expanded_depths,0,depth_axis+1)
    # Apply variable mask to depth masks
    mask = np.ma.getmask(var)
    depth_masked = np.ma.array(expanded_depths,mask=mask )
    # Look up maximum depth and surface depth to calculate total depth of water column
    max_depths = np.ma.max(depth_masked, axis=depth_axis)
    surface_depths = depth_masked.take(0,axis=depth_axis)
    total_depth = max_depths-surface_depths
    # Divide integral by total depth
    average = integral/total_depth
    
    
    return average


# # Testing

# ## Depth increasing linearly

# In[3]:


#depths (linear)
depths = np.arange(0,40)
print depths.shape


# Zero array
# 

# In[4]:


#1D zeros array
var = np.zeros((depths.shape[0],1))
print depth_average(var,depths,depth_axis=0)
print var.shape


# Testing size. What happens if the shape of the array is (N,) instead of (N,1)

# In[5]:


#1D zeros array
var = np.zeros((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape


# Constant array

# In[6]:


#1D ones array
var = np.ones((depths.shape[0],1))
print depth_average(var,depths,depth_axis=0)
print var.shape


# First axis time, second axis depth

# In[7]:


#1D ones array
#num output times
ts = 10;
var = np.ones((ts,depths.shape[0]))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape


# First axis time, second axis depth, thrid axis x

# In[8]:


#1D ones array
#num output times
ts = 10; xs=11
var = np.ones((ts,depths.shape[0],xs))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape


# First axis time, second axis depth, third axis x, fourth axis y

# In[9]:


#1D ones array
#num output times
ts = 1; xs=2; ys=3
var = np.ones((ts,depths.shape[0],xs,ys))
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape


# Move depth axis

# In[10]:


#1D ones array
#num output times
ts = 1; xs=2;
var = np.ones((ts,xs,depths.shape[0]))
avg= depth_average(var,depths,depth_axis=2)
print avg
print avg.shape
print var.shape


# Different values along time axis

# In[11]:


ts = 2; xs=3;
var = np.ones((ts,depths.shape[0],xs))
var[0,...] = 2*var[0,...]
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
print 'First time'
print avg[0,:]
print 'Second time'
print avg[1,:]


# Different values along x axis

# In[12]:


ts = 3; xs=2;
var = np.ones((ts,depths.shape[0],xs))
var[...,1] = 2*var[...,1]
avg= depth_average(var,depths,depth_axis=1)
print avg
print avg.shape
print var.shape
print 'First x'
print avg[...,0]
print 'Second x'
print avg[...,1]


# ## Depth grid spacing varies

# In[13]:


# Load depths from model data.
run='dwr_notsmooth'
base='/data/nsoontie/MEOPAR/SalishSea/results/stratification/'
path = os.path.join(base,'{}/SalishSea_1d_20030819_20030927_grid_T.nc'.format(run))
f = nc.Dataset(path,'r');
depths = f.variables['deptht']
print depths.shape


# Zeros

# In[14]:


var = np.zeros((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape


# Ones

# In[15]:


var = np.ones((depths.shape[0]))
print depth_average(var,depths,depth_axis=0)
print var.shape


# Some small error here. But why...?

# In[16]:


print depths[-1] - depths[0]


# In[17]:


diffs = depths[1:]-depths[0:-1]
s = np.sum(diffs)
print s


# In[18]:


inte =  np.trapz(var, x=depths[:],axis=0)
print inte
print inte-s


# In this case, the trapezoid rule should just give us the sum of the grid spacings. There may be some a mismatch in the floating point precision, causing the error. But the error is small and I won't worry about it.

# ## Masking

# Entire array masked, 1D

# In[19]:


var = np.zeros((depths.shape[0]))
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=0)
print var.shape
print var


# Entire array masked, multiple dimensions

# In[20]:


var = np.zeros((3,depths.shape[0]))
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var


# Part of array is masked over entire water column

# In[21]:


var = np.ones((4,depths.shape[0]))
var[0,:]=0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var


# Mask a bunch of the values

# In[23]:


var = np.ones((depths.shape[0]))
var[10:] = 0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=0)
print var.shape
print var


# Multidimensional masking

# In[24]:


ts=2; xs=3
var = np.ones((ts,depths.shape[0],xs))
var[0,10:,0] = 0
var=np.ma.masked_values(var,0)
print depth_average(var,depths,depth_axis=1)
print var.shape
print var[0,:,0]


# #Summary

# I'm pretty satisified with how this is working. There are some small rouding errors. And it must be used with maskd data. 
# 
# Next I will add this function to analyze module.

# In[ ]: