import matplotlib.pyplot as plt 
import xarray as xr
import xarray.ufuncs as uf
import numpy as np
import warnings
import gsw
import traceback
from .COAsT import COAsT
from .NEMO import NEMO
from scipy import interpolate
from scipy.integrate import cumtrapz
from sklearn.neighbors import BallTree
from skimage import measure

# =============================================================================
# The contour module is a place for code related to contours only
# =============================================================================

class Contour:
    GRAVITY = 9.8 # m s^-2
    EARTH_ROT_RATE = 7.2921 * 10**(-5) # rad/s
    
    
    @staticmethod
    def get_contours(nemo: COAsT, contour_depth: int):
        '''
        A method to obtain the continuous isbobath contours within a supplied 
        Nemo domain as a set of y indices and x indices for the model grid.

        Parameters
        ----------
        nemo : COAsT
            The Nemo object containing the dataset with the 'bathymetry' variable
        contour_depth : int
            Depth of desired contours

        Returns
        -------
        List of 2d ndarrays
            Each item of the list contains a different continuous isobath contour as a
            2d ndarray of indicies, i.e. for each list item:
            contour[:,0] contains the y indices for the contour on the model grid
            contour[:,1] contains the x indices for the contour on the model grid

        '''
        contours = measure.find_contours( nemo.dataset.bathymetry.data, contour_depth )
        # The find_contours method returns indices that have been interpolated
        # between grid points so we must round and cast to integer 
        contours = [np.round(contour).astype(int) for contour in contours]
        return contours, len(contours)
    
    
    @staticmethod
    def plot_contour(nemo: COAsT, contour: np.ndarray):
        '''
        Quick plot method to plot a contour over a pcolormesh of the 
        model bathymetry

        Parameters
        ----------
        nemo : COAsT
            The Nemo object containing the dataset with the 'bathymetry' variable
        contour : 2d array
            contour[:,0] contains the y indices for the contour on the model grid
            contour[:,1] contains the x indices for the contour on the model grid
            i.e. contour = np.vstack((y_indices,x_indices)).T

        Returns
        -------
        None

        '''
        fig,ax=plt.subplots()
        lat = nemo.dataset.latitude[xr.DataArray(contour[:,0]), xr.DataArray(contour[:,1])]
        lon = nemo.dataset.longitude[xr.DataArray(contour[:,0]), xr.DataArray(contour[:,1])]
        
        nemo.dataset.bathymetry.where(nemo.dataset.bathymetry > 0, np.nan) \
            .plot.pcolormesh(y='latitude',x='longitude',ax=ax)
        ax.scatter(lon,lat, s=0.5, color='r')
        
        
    @staticmethod    
    def get_contour_segment(nemo: COAsT, contour: np.ndarray, start_coords: np.ndarray, end_coords: np.ndarray):
        '''
        Method that will take a contour from the list of contours generated by
        coast.Contour.get_contours() and trim it to start at supplied (lat,lon)
        coordinates and end at supplied (lat, lon) coordinates.

        Parameters
        ----------
        nemo : COAsT
            The Nemo object containing the dataset with the 'bathymetry' variable
        contour : numpy.ndarray
            contour[:,0] contains the y indices for the contour on the model grid
            contour[:,1] contains the x indices for the contour on the model grid
        start_coords : numpy.ndarray
            1d array containing [latitude,longitude] of the start point of the contour
        end_coords : numpy.ndarray
            1d array containing [latitude,longitude] of the end point of the contour

        Returns
        -------
        y_ind : numpy.ndarray
            y indices of the contour on the model grid
        x_ind : numpy.ndarray
            x indices of the contour on the model grid
        contour : numpy.ndarray
            For the convenience of plotting using coast.Contour.plot_contour()
            contour[:,0] = y_ind
            contour[:,1] = x_ind

        '''

        y_ind = contour[:,0]
        x_ind = contour[:,1]
        
        # Create tree of lat and lon on the pre-processed contour
        bt = BallTree( np.deg2rad( list( zip( nemo.dataset.latitude.values[y_ind, x_ind], 
                    nemo.dataset.longitude.values[y_ind, x_ind] ) ) ), metric='haversine' )

        # Get start and end indices for contour and subset accordingly
        start_idx = bt.query( np.deg2rad( [start_coords]) )[1][0][0]
        end_idx = bt.query( np.deg2rad([end_coords]) )[1][0][0]   
        if start_idx > end_idx:
            y_ind = y_ind[end_idx:start_idx+1] 
            x_ind = x_ind[end_idx:start_idx+1]
        else:
            y_ind = y_ind[start_idx:end_idx+1] 
            x_ind = x_ind[start_idx:end_idx+1]
            
        # Ensure that the start point is closer to southern boundary of domain.
        # If start and end point have same latitude then ensure start point is 
        # closer to the western boundary of the domain.
        if y_ind[0] > y_ind[-1]:
            y_ind = y_ind[::-1]
            x_ind = x_ind[::-1]
        elif y_ind[0] == y_ind[-1]:
            if x_ind[0] > x_ind[-1]:
                y_ind = y_ind[::-1]
                x_ind = x_ind[::-1]
                
        return y_ind, x_ind, np.vstack((y_ind,x_ind)).T
                        
    
    def __init__(self, nemo: COAsT, y_ind, x_ind, depth: int):        
        '''
        Class defining a Contour type, which is a 3d dataset of points between a point A and 
        a point B defining an isobath contour. The dataset has a time, depth and contour dimension. 
        The cotnour dimension defines the points along the contour.
        The supplied model Data is subsetted in its entirety along these dimensions and
        calculations can be performed on this dataset.
        
        Parameters
        ----------
        nemo : COAsT
            nemo object containing the model dataset.
        y_ind : numpy.ndarray
            1d array of y indices defining the contour on the model grid
        x_ind : numpy.ndarray
            1d array of x indices defining the contour on the model grid
        depth : int
            Depth of contour isobath
        '''
        try:
            if y_ind[0] > y_ind[-1]:
                raise ValueError("Start point of the contour " \
                                 "must be closer than the end point of the " \
                                 "contour to the southern boundary of the model " \
                                 "domain.")
            elif y_ind[0] == y_ind[-1]:
                if x_ind[0] > x_ind[-1]:
                    raise ValueError("Start and end points of the contour " \
                                     "have the same latitudes, the start point must " \
                                     "be the closer of the two points to the western " \
                                     "boundary of the model domain.")

        
            self.depth = depth
            self.y_ind, self.x_ind = self.process_contour( nemo.dataset, y_ind, x_ind )
            self.len = len(self.y_ind)
            self.filename_domain = nemo.filename_domain
            da_y_ind = xr.DataArray( self.y_ind, dims=['r_dim'] )
            da_x_ind = xr.DataArray( self.x_ind, dims=['r_dim'] )
            self.data_contour = nemo.dataset.isel(y_dim = da_y_ind, x_dim = da_x_ind)
        except ValueError:
            print(traceback.format_exc())
        
        
    def process_contour(self, ds: xr.Dataset, y_ind, x_ind):
        
        """ Redefine contour so that each point on the contour defined by
        y_ind and x_ind is seperated from its neighbours by a single index change 
        in y or x, but not both. 
        example: convert y_ind = [10,11], x_ind = [1,2] to y_ind = [10,10], x_ind = [1,2]
        or  y_ind = [10,11], x_ind = [1,1]
         
        Parameters
        ----------
        nemo : xarray.Dataset
            xarray Dataset from supplied nemo object
        y_ind : numpy.ndarray
            1d array of y indices defining the contour on the model grid
        x_ind : numpy.ndarray
            1d array of x indices defining the contour on the model grid

        Returns
        -------
        y_ind : numpy.ndarray
            processed y indices of the contour on the model grid
        x_ind : numpy.ndarray
            processed x indices of the contour on the model grid
         
        """
         
        try:
            y_ind = np.asarray(y_ind)
            x_ind = np.asarray(x_ind)
            # When replacing diagonal segments in the contour, pick the path that is
            # closest to the contour isobath depth
            option1 = uf.fabs(ds.bathymetry[xr.DataArray(y_ind+1), xr.DataArray(x_ind)] - self.depth)
            option0 = uf.fabs(ds.bathymetry[xr.DataArray(y_ind), xr.DataArray(x_ind+1)] - self.depth)
            add_new_y_point = xr.where(option1 <= option0, 1, 0 )
        
            spacing = np.abs( np.diff(y_ind) ) + np.abs( np.diff(x_ind) )
            if spacing.max() > 2:
                raise ValueError("The contour is not continuous. The contour must be defined on " 
                                 "adjacent grid points.")
            spacing[spacing!=2]=0
            doublespacing = np.nonzero( spacing )[0]
            for ispacing in doublespacing[::-1]:
                if add_new_y_point[ispacing]:
                    y_ind = np.insert( y_ind, ispacing+1, y_ind[ispacing+1] )
                    x_ind = np.insert( x_ind, ispacing+1, x_ind[ispacing] )
                else:
                    y_ind = np.insert( y_ind, ispacing+1, y_ind[ispacing] )
                    x_ind = np.insert( x_ind, ispacing+1, x_ind[ispacing+1] ) 
            
            # Remove any repeated points caused by the rounding of the indices
            nonrepeated_idx = np.nonzero( np.abs( np.diff(y_ind) ) + np.abs( np.diff(x_ind) ) )
    
            y_ind = np.concatenate( (y_ind[nonrepeated_idx], [y_ind[-1]]) )
            x_ind = np.concatenate( (x_ind[nonrepeated_idx], [x_ind[-1]]) )
    
            return (y_ind, x_ind)
        except ValueError:
            print(traceback.format_exc())
   
    
    def gen_z_levels(self, max_depth):
        ''' Generates a pre-defined 1d vertical depth coordinates,
        i.e. horizontal z-level vertical coordinates up to a supplied 
        maximum depth, 'max_depth' '''
        
        max_depth = max_depth + 650
        z_levels_0_50 = np.arange(0,55,5)
        z_levels_60_290 = np.arange(60,300,10)
        z_levels_300_600 = np.arange(300,650,50)
        z_levels_650_ = np.arange(650,max_depth+150,150)
        z_levels = np.concatenate( (z_levels_0_50, z_levels_60_290, 
                                    z_levels_300_600, z_levels_650_) )
        z_levels = z_levels[z_levels <= max_depth] 
        return z_levels
    

class Contour_f(Contour):
    '''
    Class defining a Contour type on the f-grid, which is a 3d dataset of points between a point A and 
    a point B defining an isobath contour. The dataset has a time, depth and contour dimension. 
    The contour dimension defines the points along the contour.
    The supplied model f-grid Data is subsetted in its entirety along these dimensions 
    within Contour_f.data_contour of type xarray.Dataset

    
    Parameters
    ----------
    nemo : COAsT
        f-grid nemo object containing the model dataset.
    y_ind : numpy.ndarray
        1d array of y indices defining the contour on the model grid
    x_ind : numpy.ndarray
        1d array of x indices defining the contour on the model grid
    depth : int
        Depth of contour isobath
    '''
    
    def __init__(self, nemo_f: COAsT, y_ind, x_ind, depth):
        super().__init__(nemo_f, y_ind, x_ind, depth)
        self.data_cross_flow = xr.Dataset()
     
        
    def calc_cross_contour_flow(self, nemo_u: COAsT, nemo_v: COAsT):
        '''
        Method that will calculate the flow across the contour and store this data
        within Contour_f.data_cross_flow, which is an xarray.Dataset. Specifically
        Contour_f.normal_velocities are the velocities across the contour
        (time, depth, position along contour) in m/s
        Contour_f.depth_integrated_normal_transport are the depth integrated 
        volume transports across the contour (time, position along contour) in Sv

        Parameters
        ----------
        nemo_u : COAsT
            The nemo object containing the model data on the u-grid.
        nemo_v : COAsT
            The nemo object containing the model data on the v-grid.

        Returns
        -------
        None.

        '''
                
        # subset the u and v datasets 
        da_y_ind = xr.DataArray( self.y_ind, dims=['r_dim'] )
        da_x_ind = xr.DataArray( self.x_ind, dims=['r_dim'] )
        u_ds = nemo_u.dataset.isel(y_dim = da_y_ind, x_dim = da_x_ind)
        v_ds = nemo_v.dataset.isel(y_dim = da_y_ind, x_dim = da_x_ind)
        
        dr_n = np.where(np.diff(self.y_ind)>0, np.arange(0,u_ds.r_dim.size-1), np.nan )
        dr_n = dr_n[~np.isnan(dr_n)].astype(int)
        dr_s = np.where(np.diff(self.y_ind)<0, np.arange(0,u_ds.r_dim.size-1), np.nan )
        dr_s = dr_s[~np.isnan(dr_s)].astype(int)
        dr_e = np.where(np.diff(self.x_ind)>0, np.arange(0,v_ds.r_dim.size-1), np.nan )
        dr_e = dr_e[~np.isnan(dr_e)].astype(int)
        dr_w = np.where(np.diff(self.x_ind)<0, np.arange(0,v_ds.r_dim.size-1), np.nan )
        dr_w = dr_w[~np.isnan(dr_w)].astype(int)
        
        # Note that subsetting the dataset first instead of subsetting each array seperately,
        # as we do here, is neater but significantly slower.
        self.data_cross_flow['normal_velocities'] = xr.full_like(u_ds.vozocrtx, np.nan)        
        self.data_cross_flow['normal_velocities'][:,:,dr_n] = u_ds.vozocrtx.data[:,:,dr_n+1]
        self.data_cross_flow['normal_velocities'][:,:,dr_s] = -u_ds.vozocrtx.data[:,:,dr_s]
        self.data_cross_flow['normal_velocities'][:,:,dr_e] = -v_ds.vomecrty.data[:,:,dr_e+1]
        self.data_cross_flow['normal_velocities'][:,:,dr_w] = v_ds.vomecrty.data[:,:,dr_w]
        self.data_cross_flow['normal_velocities'].attrs = {'units':'m/s', \
                'standard_name':'contour-normal velocities'}
             
        self.data_cross_flow['normal_transport'] = xr.full_like(u_ds.vozocrtx, np.nan)  
        self.data_cross_flow['normal_transport'][:,:,dr_n] = ( u_ds.vozocrtx.data[:,:,dr_n+1] * 
                                u_ds.e2.data[dr_n+1] * u_ds.e3_0.data[:,dr_n+1] )
        self.data_cross_flow['normal_transport'][:,:,dr_s] = ( -u_ds.vozocrtx.data[:,:,dr_s] * 
                                u_ds.e2.data[dr_s] * u_ds.e3_0.data[:,dr_s] )
        self.data_cross_flow['normal_transport'][:,:,dr_e] = ( -v_ds.vomecrty.data[:,:,dr_e+1] *
                                v_ds.e1.data[dr_e+1] * v_ds.e3_0.data[:,dr_e+1] )
        self.data_cross_flow['normal_transport'][:,:,dr_w] = ( v_ds.vomecrty.data[:,:,dr_w] *
                                v_ds.e1.data[dr_w] * v_ds.e3_0.data[:,dr_w] )
        self.data_cross_flow['normal_transport'].attrs = {'units':'m^3/s', \
                'standard_name':'contour-normal volume transport'}
        
        self.data_cross_flow['depth_integrated_normal_transport'] = (self.data_cross_flow
                                .normal_transport.sum(dim='z_dim') / 1000000.)
        self.data_cross_flow['depth_integrated_normal_transport'].attrs ={'units':'Sv', \
                'standard_name':'contour-normal depth integrated volume transport'}
                                
        self.__update_cross_flow_vars('depth_0',u_ds.depth_0,v_ds.depth_0,dr_n,dr_s,dr_e,dr_w,1)
        self.__update_cross_flow_vars('longitude',u_ds.longitude,v_ds.longitude,dr_n,dr_s,dr_e,dr_w,0)
        self.__update_cross_flow_vars('latitude',u_ds.latitude,v_ds.latitude,dr_n,dr_s,dr_e,dr_w,0)
        self.data_cross_flow['e1'] = xr.full_like(self.data_contour.e1, np.nan)   
        self.__update_cross_flow_vars('e1',u_ds.e1,v_ds.e1,dr_n,dr_s,dr_e,dr_w,0)
        self.data_cross_flow['e2'] = xr.full_like(self.data_contour.e2, np.nan)
        self.__update_cross_flow_vars('e2',u_ds.e2,v_ds.e2,dr_n,dr_s,dr_e,dr_w,0)
        self.data_cross_flow['e3_0'] = xr.full_like(self.data_contour.e3_0, np.nan)
        self.__update_cross_flow_vars('e3_0',u_ds.e3_0,v_ds.e3_0,dr_n,dr_s,dr_e,dr_w,1)
        
        self.data_cross_flow['depth_0'].attrs = {'standard_name':'Depth at time zero \
                on the contour-normal velocity grid points'}
        self.data_cross_flow['latitude'].attrs = {'standard_name':'Latitude at \
                the contour-normal velocity grid points'}
        self.data_cross_flow['longitude'].attrs = {'standard_name':'Longitude at \
                the contour-normal velocity grid points'}

    
    def __update_cross_flow_vars(self, var, u_var,v_var, dr_n, dr_s, dr_e, dr_w, pos ):
        ''' This method will pull variable data at specific points along the contour
        from the u and v grid datasets and put them into the data_cross_flow dataset'''
        if pos==0:
            self.data_cross_flow[var][dr_n] = u_var.data[dr_n+1]
            self.data_cross_flow[var][dr_s] = u_var.data[dr_s]
            self.data_cross_flow[var][dr_e] = v_var.data[dr_e+1]
            self.data_cross_flow[var][dr_w] = v_var.data[dr_w]
            self.data_cross_flow[var][-1] = np.nan  
        elif pos==1:         
            self.data_cross_flow[var][:,dr_n] = u_var.data[:,dr_n+1]
            self.data_cross_flow[var][:,dr_s] = u_var.data[:,dr_s]
            self.data_cross_flow[var][:,dr_e] = v_var.data[:,dr_e+1]
            self.data_cross_flow[var][:,dr_w] = v_var.data[:,dr_w]
            self.data_cross_flow[var][:,-1] = np.nan    
        elif pos==2:         
            self.data_cross_flow[var][:,:,dr_n] = u_var.data[:,:,dr_n+1]
            self.data_cross_flow[var][:,:,dr_s] = u_var.data[:,:,dr_s]
            self.data_cross_flow[var][:,:,dr_e] = v_var.data[:,:,dr_e+1]
            self.data_cross_flow[var][:,:,dr_w] = v_var.data[:,:,dr_w]
            self.data_cross_flow[var][:,:,-1] = np.nan    
        
            
    def __pressure_grad_fpoint2(self, ds_T, ds_T_j1, ds_T_i1, ds_T_j1i1, r_ind, velocity_component):
        """
        Calculates the hydrostatic and surface pressure gradients at a set of f-points
        along the contour, i.e. at a set of specific values of r_dim (but for all time and depth).
        The caller must supply four datasets that contain the variables which define
        the hydrostatic and surface pressure at all vertical z_levels and all time 
        on the t-points around the contour i.e. for a set of f-points on the contour 
        defined each defined at (j+1/2, i+1/2), we want t-points at (j,i), (j+1,i), (j,i+1), (j+1,i+1), 
        corresponding to ds_T, ds_T_j1, ds_T_i1, ds_T_j1i1, respectively. 
        ds_T, ds_T_j1, ds_T_i1, ds_T_j1i1 will have dimensions in time and depth.
        
        The velocity_component defines whether u or v is normal to the contour 
        for the segments of the contour. A segment of contour is 
        defined as being r_dim to r_dim+1 where r_dim is the along contour dimension.


        Returns
        -------
        hpg_f : DataArray with dimensions in time and depth and along contour
            hydrostatic pressure gradient at a set of f-points along the contour
            for all time and depth
        spg_f : DataArray with dimensions in time and depth and along contour
            surface pressure gradient at a set of f-points along the contour

        """
        if velocity_component == "u":
            # required scale factors for derivative and averaging
            e2v = 0.5*( ds_T_j1.e2.data[r_ind] + ds_T.e2.data[r_ind] )
            e2v_i1 = 0.5*( ds_T_j1i1.e2.data[r_ind] + ds_T_i1.e2.data[r_ind] )
            e1v = 0.5*( ds_T_j1.e1.data[r_ind] + ds_T.e1.data[r_ind] )
            e1v_i1 = 0.5*( ds_T_j1i1.e1.data[r_ind] + ds_T_i1.e1.data[r_ind] )
            e1f = 0.5*( e1v + e1v_i1 )            
            # calculate gradients at v-points either side of f-point
            hpg = (ds_T_j1.pressure_h_zlevels.data[:,:,r_ind] - ds_T.pressure_h_zlevels.data[:,:,r_ind]) / e2v
            hpg_i1 = (ds_T_j1i1.pressure_h_zlevels.data[:,:,r_ind] - ds_T_i1.pressure_h_zlevels.data[:,:,r_ind]) / e2v_i1   
            # average onto f-point
            hpg_f = 0.5 * ( ( e1v * hpg ) + ( e1v_i1 * hpg_i1 ) ) / e1f 
            # as aboave            
            spg = (ds_T_j1.pressure_s.data[:,r_ind] - ds_T.pressure_s.data[:,r_ind]) / e2v
            spg_i1 = (ds_T_j1i1.pressure_s.data[:,r_ind] - ds_T_i1.pressure_s.data[:,r_ind]) / e2v_i1
            spg_f = 0.5 * ( (e1v * spg) + (e1v_i1 * spg_i1) ) / e1f 
        elif velocity_component == "v":
            # required scale factors for derivative and averaging
            e1u = 0.5 * ( ds_T_i1.e1.data[r_ind] + ds_T.e1.data[r_ind] ) 
            e1u_j1 = 0.5 * ( ds_T_j1i1.e1.data[r_ind] + ds_T_j1.e1.data[r_ind] )
            e2u = 0.5 * ( ds_T_i1.e2.data[r_ind] + ds_T.e2.data[r_ind] )
            e2u_j1 = 0.5 * ( ds_T_j1i1.e2.data[r_ind] + ds_T_j1.e2.data[r_ind] )
            e2f = 0.5 * ( e2u + e2u_j1 )
            # calculate gradients at u-points either side of f-point
            hpg = (ds_T_i1.pressure_h_zlevels.data[:,:,r_ind] - ds_T.pressure_h_zlevels.data[:,:,r_ind]) / e1u
            hpg_j1 = (ds_T_j1i1.pressure_h_zlevels.data[:,:,r_ind] - ds_T_j1.pressure_h_zlevels.data[:,:,r_ind]) / e1u_j1 
            # average onto f-point
            hpg_f = 0.5 * ( (e2u * hpg) + (e2u_j1 * hpg_j1) ) / e2f
            # as above
            spg = (ds_T_i1.pressure_s.data[:,r_ind] - ds_T.pressure_s.data[:,r_ind]) / e1u
            spg_j1 = (ds_T_j1i1.pressure_s.data[:,r_ind] - ds_T_j1.pressure_s.data[:,r_ind]) / e1u_j1 
            spg_f = 0.5 * ( (e2u * spg) + (e2u_j1 * spg_j1) ) / e2f
        
        return (hpg_f, spg_f)
   
    
    def calc_geostrophic_flow(self, nemo_t: COAsT, ref_density=None):
        """
        This method will calculate the geostrophic velocity and volume transport
        (due to the geostrophic current) across the contour. 
        Four variables are added to the Contour.data_cross_flow dataset:
        1. normal_velocity_hpg      (t_dim, depth_z_levels, r_dim)
        This is the velocity due to the hydrostatic pressure gradient
        2. normal_velocity_spg      (t_dim, r_dim)
        This is the velocity due to the surface pressure gradient
        3. transport_across_AB_hpg  (t_dim, r_dim)
        This is the volume transport due to the hydrostatic pressure gradient
        4. transport_across_AB_spg  (t_dim, r_dim
        This is the volume transport due to the surface pressure gradient
                                                                       
        This implementation works by regridding vertically onto horizontal z_levels in order
        to perform the horizontal gradients. Currently s_level depths are
        assumed fixed at their initial depths, i.e. at time zero.
        
        Requirements: The nemo t-grid dataset, nemo_t, must contain the sea surface height,
        Practical Salinity and the Potential Temperature variables. The depth_0
        field must also be supplied. The GSW package is used to calculate
        The Absolute Pressure, Absolute Salinity and Conservate Temperature.
        
        Parameters
        ----------
        nemo_t : COAsT
            This is the nemo model data on the t-grid for the entire domain.
        ref_density : TYPE, optional
            reference density value. If not supplied a mean in time, depth and 
            along the contour will be used as the mean reference value.

        Returns
        -------
        None.

        """
        # If there is no time dimension, add one then remove at end. This is so
        # indexing can assume a time dimension exists
        nemo_t_local = nemo_t.copy()
        if 't_dim' not in nemo_t_local.dataset.dims:
            nemo_t_local.dataset = nemo_t_local.dataset.expand_dims(dim={'t_dim':1},axis=0)
        
        # We need to calculate the pressure at four t-points to get an
        # average onto the pressure gradient at the f-points, which will then
        # be averaged onto the normal velocity points. Here we subset the nemo_t 
        # data around the contour so we have these four t-grid points at each 
        # point along the contour        
        cont_t = Contour_t(nemo_t_local, self.y_ind, self.x_ind, self.depth)            # j,i
        cont_t_j1 = Contour_t(nemo_t_local, self.y_ind+1, self.x_ind, self.depth)       # j+1,i
        cont_t_i1 = Contour_t(nemo_t_local, self.y_ind, self.x_ind+1, self.depth)       # j,i+1
        cont_t_j1i1 = Contour_t(nemo_t_local, self.y_ind+1, self.x_ind+1, self.depth)   # j+1,i+1
        
        bath_max = np.max([cont_t.data_contour.bathymetry.max().item(), 
                           cont_t_j1.data_contour.bathymetry.max().item(),
                           cont_t_i1.data_contour.bathymetry.max().item(), 
                           cont_t_j1i1.data_contour.bathymetry.max().item()])   
        z_levels = self.gen_z_levels(bath_max)
        
        cont_t.construct_pressure(ref_density, z_levels, extrapolate=True)
        cont_t_j1.construct_pressure(ref_density, z_levels, extrapolate=True)
        cont_t_i1.construct_pressure(ref_density, z_levels, extrapolate=True)
        cont_t_j1i1.construct_pressure(ref_density, z_levels, extrapolate=True)        
        
        # Remove the mean hydrostatic pressure on each z_level from the hydrostatic pressure.
        # This helps to reduce the noise when taking the horizontal gradients of hydrostatic pressure.
        # Also catch and ignore nan-slice warning
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=RuntimeWarning)
            pressure_h_zlevel_mean = ( xr.concat( (cont_t.data_contour.pressure_h_zlevels, 
                            cont_t_j1.data_contour.pressure_h_zlevels, 
                            cont_t_i1.data_contour.pressure_h_zlevels, 
                            cont_t_j1i1.data_contour.pressure_h_zlevels), dim='concat_dim' )
                            .mean(dim=('concat_dim','r_dim','t_dim'),skipna=True) )
            if ref_density is None:
                ref_density = ( xr.concat( (cont_t.data_contour.density_zlevels, 
                            cont_t_j1.data_contour.density_zlevels, 
                            cont_t_i1.data_contour.density_zlevels, 
                            cont_t_j1i1.data_contour.density_zlevels), dim='concat_dim' )
                            .mean(dim=('concat_dim','r_dim','t_dim','depth_z_levels'),skipna=True).item() )
        cont_t.data_contour['pressure_h_zlevels'] = \
                cont_t.data_contour.pressure_h_zlevels - pressure_h_zlevel_mean
        cont_t_j1.data_contour['pressure_h_zlevels'] = \
                cont_t_j1.data_contour.pressure_h_zlevels - pressure_h_zlevel_mean
        cont_t_i1.data_contour['pressure_h_zlevels'] = \
                cont_t_i1.data_contour.pressure_h_zlevels - pressure_h_zlevel_mean
        cont_t_j1i1.data_contour['pressure_h_zlevels'] = \
                cont_t_j1i1.data_contour.pressure_h_zlevels - pressure_h_zlevel_mean
                
        # Coriolis parameter
        f = 2 * self.EARTH_ROT_RATE * np.sin( np.deg2rad(self.data_contour.latitude) )        
    
        # Find the indices where the derivative of the contour in the north, south, east and west
        # directions are positive.
        dr_n = np.where(np.diff(self.y_ind)>0, np.arange(0,self.data_contour.r_dim.size-1), np.nan )
        dr_s = np.where(np.diff(self.y_ind)<0, np.arange(0,self.data_contour.r_dim.size-1), np.nan )
        dr_e = np.where(np.diff(self.x_ind)>0, np.arange(0,self.data_contour.r_dim.size-1), np.nan )
        dr_w = np.where(np.diff(self.x_ind)<0, np.arange(0,self.data_contour.r_dim.size-1), np.nan )        
        dr_list = [dr_n[~np.isnan(dr_n)].astype(int), dr_s[~np.isnan(dr_s)].astype(int),
                   dr_e[~np.isnan(dr_e)].astype(int), dr_w[~np.isnan(dr_w)].astype(int)]
        
        # horizontal scale factors on the relevent u and v grids that are
        # normal to the contour for dr_n, dr_s, dr_e, dr_w
        e2u_j1  = 0.5 * ( cont_t_j1.data_contour.e2.data[dr_list[0]] + cont_t_j1i1.data_contour.e2.data[dr_list[0]] )
        e2u     = 0.5 * ( cont_t.data_contour.e2.data[dr_list[1]] + cont_t_i1.data_contour.e2.data[dr_list[1]] )
        e1v_i1  = 0.5 * ( cont_t_i1.data_contour.e1.data[dr_list[2]] + cont_t_j1i1.data_contour.e1.data[dr_list[2]] )
        e1v     = 0.5 * ( cont_t.data_contour.e1.data[dr_list[3]] + cont_t_j1.data_contour.e1.data[dr_list[3]] )
        e_horiz_vel = [e2u_j1, e2u, e1v_i1, e1v] 
        # Horizontal scale factors on f-grid for dr_n, dr_s, dr_e, dr_w
        e_horiz_f   = [self.data_contour.e2, self.data_contour.e2, 
                       self.data_contour.e1, self.data_contour.e1]
        # velocity component normal to contour for dr_n, dr_s, dr_e, dr_w
        velocity_component = ["u","u","v",'v']
        # Geostrophic flow direction across contour
        flow_direction = [-1,1,-1,1]        
        
        normal_velocity_hpg = np.zeros_like(cont_t.data_contour.pressure_h_zlevels)
        normal_velocity_spg = np.zeros_like(cont_t.data_contour.pressure_s)
        # horizontal scale factors for each segmant of contour
        e_horiz = np.zeros( (cont_t.data_contour.t_dim.size, cont_t.data_contour.r_dim.size) ) 
        # Contruct geostrophic flow
        for dr, vel_comp, flow_dir, e_hor_vel, e_hor_f in \
                zip(dr_list, velocity_component, flow_direction, e_horiz_vel, e_horiz_f) :
            hpg, spg        = self.__pressure_grad_fpoint2( cont_t.data_contour, 
                                cont_t_j1.data_contour, cont_t_i1.data_contour, 
                                cont_t_j1i1.data_contour, dr, vel_comp )
            hpg_r1, spg_r1  = self.__pressure_grad_fpoint2( cont_t.data_contour,
                                cont_t_j1.data_contour,
                                cont_t_i1.data_contour, 
                                cont_t_j1i1.data_contour, dr+1, vel_comp )
            normal_velocity_hpg[:,:,dr] = ( flow_dir * 0.5 * (e_hor_f.data[dr]*hpg/f.data[dr] 
                                            + e_hor_f.data[dr+1]*hpg_r1/f.data[dr+1]) 
                                            / (e_hor_vel * ref_density) )
            normal_velocity_spg[:,dr]   = ( flow_dir * 0.5 * (e_hor_f.data[dr]*spg/f.data[dr] 
                                            + e_hor_f.data[dr+1]*spg_r1/f.data[dr+1]) 
                                            / (e_hor_vel * ref_density) ) 
            e_horiz[:,dr]               = e_hor_vel

        # Bathymetry at normal velocity points             
        H = np.zeros_like( self.data_contour.bathymetry.values )
        H[:-1] = 0.5*(self.data_contour.bathymetry.values[:-1] + self.data_contour.bathymetry.values[1:])
        # Remove redundent levels below bathymetry
        normal_velocity_hpg = np.where( z_levels[:,np.newaxis] <= H, 
                                       normal_velocity_hpg, np.nan )           
        active_z_levels = np.count_nonzero(~np.isnan(normal_velocity_hpg),axis=1).max() 
        normal_velocity_hpg = normal_velocity_hpg[:,:active_z_levels,:]
        z_levels = z_levels[:active_z_levels]

        # The cross contour flow is defined on the u and v points that are across
        # the contour, i.e. between f points, therefore the attributes of the
        # data_cross_flow dataset need to be on these points.
        da_y_ind = xr.DataArray( self.y_ind, dims=['r_dim'] )
        da_x_ind = xr.DataArray( self.x_ind, dims=['r_dim'] )
        u_ds = NEMO( fn_domain=self.filename_domain, grid_ref='u-grid' ).dataset \
            .isel(y_dim=da_y_ind, x_dim=da_x_ind)
        v_ds = NEMO( fn_domain=self.filename_domain, grid_ref='v-grid' ).dataset \
            .isel(y_dim=da_y_ind, x_dim=da_x_ind)
        self.data_cross_flow['e1'] = xr.full_like(self.data_contour.e1, np.nan)  
        self.data_cross_flow['e2'] = xr.full_like(self.data_contour.e2, np.nan)
        self.data_cross_flow['latitude'] = xr.full_like(self.data_contour.latitude, np.nan)  
        self.data_cross_flow['longitude'] = xr.full_like(self.data_contour.longitude, np.nan)
        self.__update_cross_flow_vars('longitude',u_ds.longitude,v_ds.longitude,
                            dr_list[0],dr_list[1],dr_list[2],dr_list[3],0)
        self.__update_cross_flow_vars('latitude',u_ds.latitude,v_ds.latitude,
                            dr_list[0],dr_list[1],dr_list[2],dr_list[3],0)         
        self.__update_cross_flow_vars('e1',u_ds.e1,v_ds.e1,
                            dr_list[0],dr_list[1],dr_list[2],dr_list[3],0)        
        self.__update_cross_flow_vars('e2',u_ds.e2,v_ds.e2,
                            dr_list[0],dr_list[1],dr_list[2],dr_list[3],0)
        self.data_cross_flow['latitude'].attrs = {'standard_name':'Latitude at \
                the contour-normal velocity grid points'}
        self.data_cross_flow['longitude'].attrs = {'standard_name':'Longitude at \
                the contour-normal velocity grid points'}
        
        # DataArray attributes
        coords_hpg={'depth_z_levels': (('depth_z_levels'), z_levels),
                'latitude': (('r_dim'), self.data_cross_flow.latitude),
                'longitude': (('r_dim'), self.data_cross_flow.longitude)}
        dims_hpg=['depth_z_levels', 'r_dim']
        attributes_hpg = {'units': 'm/s', 'standard name': 'velocity across the \
                          transect due to the hydrostatic pressure gradient'}
        coords_spg={'latitude': (('r_dim'), self.data_cross_flow.latitude),
                'longitude': (('r_dim'), self.data_cross_flow.longitude)}
        dims_spg=['r_dim']
        attributes_spg = {'units': 'm/s', 'standard name': 'velocity across the \
                          transect due to the surface pressure gradient'}
        
        # Add time if required
        if 't_dim' in cont_t.data_contour.dims:
            coords_hpg['time'] = (('t_dim'), cont_t.data_contour.time)
            dims_hpg.insert(0, 't_dim')
            coords_spg['time'] = (('t_dim'), cont_t.data_contour.time)
            dims_spg.insert(0, 't_dim')
        
        # Add DataArrays  to dataset
        self.data_cross_flow['normal_velocity_hpg'] = xr.DataArray( np.squeeze(normal_velocity_hpg),
                coords=coords_hpg, dims=dims_hpg, attrs=attributes_hpg)
        self.data_cross_flow['normal_velocity_spg'] = xr.DataArray( np.squeeze(normal_velocity_spg),
                coords=coords_spg, dims=dims_spg, attrs=attributes_spg)
        self.data_cross_flow['transport_across_AB_hpg'] = ( self.data_cross_flow
                .normal_velocity_hpg.fillna(0).integrate(dim='depth_z_levels') ) * e_horiz / 1000000   
        self.data_cross_flow.transport_across_AB_hpg.attrs = {'units': 'Sv', 
                'standard_name': 'volume transport across transect due to the hydrostatic pressure gradient'}        
        self.data_cross_flow['transport_across_AB_spg'] = self.data_cross_flow.normal_velocity_spg * H * e_horiz / 1000000
        self.data_cross_flow.transport_across_AB_spg.attrs = {'units': 'Sv', 
                'standard_name': 'volume transport across transect due to the surface pressure gradient'}
        
        
class Contour_t(Contour):
    '''
    Class defining a Contour type on the t-grid, which is a 3d dataset of points between a point A and 
    a point B defining an isobath contour. The dataset has a time, depth and contour dimension. 
    The contour dimension defines the points along the contour.
    The supplied model t-grid Data is subsetted in its entirety along these dimensions and
    calculations can be performed on this dataset.
    
    Parameters
    ----------
    nemo : COAsT
        t-grid nemo object containing the model dataset.
    y_ind : numpy.ndarray
        1d array of y indices defining the contour on the model grid
    x_ind : numpy.ndarray
        1d array of x indices defining the contour on the model grid
    depth : int
        Depth of contour isobath
    '''
    
    def __init__(self, nemo_t: COAsT, y_ind, x_ind, depth):
        super().__init__(nemo_t, y_ind, x_ind, depth)
        
        
    def construct_pressure( self, ref_density=None, z_levels=None, extrapolate=False ):   
        '''
            This method is for calculating the hydrostatic and surface pressure fields
            on horizontal levels in the vertical (z-levels). The motivation 
            is to enable the calculation of horizontal gradients; however, 
            the variables can quite easily be interpolated onto the original 
            vertical grid.
             
            Requirements: The object's t-grid dataset must contain the sea surface height,
            Practical Salinity and the Potential Temperature variables.
            The GSW package is used to calculate the Absolute Pressure, 
            Absolute Salinity and Conservate Temperature.
            
            Three new variables (density, hydrostatic pressure, surface pressure)
            are created and added to the Contour_t.data_contour dataset:
                density_zlevels       (t_dim, depth_z_levels, r_dim)
                pressure_h_zlevels    (t_dim, depth_z_levels, r_dim)
                pressure_s            (t_dim, r_dim)
            
            Note that density is constructed using the EOS10
            equation of state.

        Parameters
        ----------
        ref_density: float
            reference density value, if None, then the Contour mean across time, 
            depth and along contour will be used.
        z_levels : (optional) numpy array
            1d array that defines the depths to interpolate the density and pressure
            on to.
        extrapolate : boolean, default False
            If true the variables are extrapolated to the deepest z_level, if false,
            values below the bathymetry are set to NaN
        Returns
        -------
        None.

        '''        
        
        # If there is no time dimension, add one, this is so
        # indexing can assume a time dimension exists
        if 't_dim' not in self.data_contour.dims:
            self.data_contour = self.data_contour.expand_dims(dim={'t_dim':1},axis=0)

        # Generate vertical levels if not supplied
        if z_levels is None:   
            z_levels = self.gen_z_levels( self.data_contour.bathymetry.max().item() )
        
        shape_ds = ( self.data_contour.t_dim.size, len(z_levels), self.data_contour.r_dim.size )
        salinity_z = np.ma.zeros( shape_ds )
        temperature_z = np.ma.zeros( shape_ds ) 
        salinity_s = self.data_contour.salinity.to_masked_array()
        temperature_s = self.data_contour.temperature.to_masked_array()
        s_levels = self.data_contour.depth_0.values
        
        # Interpolate salinity and temperature onto z-levels
        # Note. At the current time there does not appear to be a good algorithm for 
        # performing this type of interpolation without loops, which can be a bottleneck.
        # Griddata is an option but does not support extrapolation and did not 
        # have noticable performance benefit.
        for it in self.data_contour.t_dim:
            for ir in self.data_contour.r_dim:
                if not np.all(np.isnan(salinity_s[it,:,ir].data)):  
                    # Need to remove the levels below the (envelope) bathymetry which are NaN
                    salinity_s_r = salinity_s[it,:,ir].compressed()
                    temperature_s_r = temperature_s[it,:,ir].compressed()
                    s_levels_r = s_levels[:len(salinity_s_r),ir]
                    
                    sal_func = interpolate.interp1d( s_levels_r, salinity_s_r, 
                                 kind='linear', fill_value="extrapolate")
                    temp_func = interpolate.interp1d( s_levels_r, temperature_s_r, 
                                 kind='linear', fill_value="extrapolate")
                    
                    if extrapolate is True:
                        salinity_z[it,:,ir] = sal_func(z_levels)
                        temperature_z[it,:,ir] = temp_func(z_levels)                        
                    else:
                        # set levels below the bathymetry to nan
                        salinity_z[it,:,ir] = np.where( z_levels <= self.data_contour.bathymetry.values[ir], 
                                sal_func(z_levels), np.nan )
                        temperature_z[it,:,ir] = np.where( z_levels <= self.data_contour.bathymetry.values[ir], 
                                temp_func(z_levels), np.nan ) 
                    
        if extrapolate is False:
            # remove redundent levels    
            active_z_levels = np.count_nonzero(~np.isnan(salinity_z),axis=1).max() 
            salinity_z = salinity_z[:,:active_z_levels,:]
            temperature_z = temperature_z[:,:active_z_levels,:]
            z_levels = z_levels[:active_z_levels]
        
        # Absolute Pressure (depth must be negative)   
        pressure_absolute = np.ma.masked_invalid(
            gsw.p_from_z( -z_levels[:,np.newaxis], self.data_contour.latitude ) )         
        # Absolute Salinity           
        salinity_absolute = np.ma.masked_invalid(
            gsw.SA_from_SP( salinity_z, pressure_absolute, self.data_contour.longitude, 
            self.data_contour.latitude ) )
        salinity_absolute = np.ma.masked_less(salinity_absolute,0)
        # Conservative Temperature
        temp_conservative = np.ma.masked_invalid(
            gsw.CT_from_pt( salinity_absolute, temperature_z ) )
        # In-situ density
        density_z = np.ma.masked_invalid( gsw.rho( 
            salinity_absolute, temp_conservative, pressure_absolute ) )
        
        coords={'depth_z_levels': (('depth_z_levels'), z_levels),
                'latitude': (('r_dim'), self.data_contour.latitude),
                'longitude': (('r_dim'), self.data_contour.longitude)}
        dims=['depth_z_levels', 'r_dim']
        attributes = {'units': 'kg / m^3', 'standard name': 'In-situ density on the z-level vertical grid'}
        
        if shape_ds[0] != 1:
            coords['time'] = (('t_dim'), self.data_contour.time.values)
            dims.insert(0, 't_dim')
         
        if ref_density is None:    
            ref_density = np.mean( density_z )
        self.data_contour['density_zlevels'] = xr.DataArray( np.squeeze(density_z), 
                coords=coords, dims=dims, attrs=attributes )
    
        # Cumulative integral of perturbation density on z levels
        density_cumulative = -cumtrapz( density_z - ref_density, x=-z_levels, axis=1, initial=0)
        hydrostatic_pressure = density_cumulative * self.GRAVITY
        
        attributes = {'units': 'kg m^{-1} s^{-2}', 'standard name': 'Hydrostatic perturbation pressure on the z-level vertical grid'}
        self.data_contour['pressure_h_zlevels'] = xr.DataArray( np.squeeze(hydrostatic_pressure), 
                coords=coords, dims=dims, attrs=attributes )        
        self.data_contour['pressure_s'] = ref_density * self.GRAVITY * self.data_contour.ssh.squeeze()
        self.data_contour.pressure_s.attrs = {'units': 'kg m^{-1} s^{-2}', 
                                  'standard_name': 'Surface perturbation pressure'}
        
    
        
        
    