#!/usr/bin/env python

from __future__ import print_function

import sys
import argparse
import datetime
import netCDF4 as nc
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

"""
What this script does:

Generate a plot of Atlantic and Global Meridional Overturning Circulation.

How to run this script:

./mom_plot_moc.py ocean.nc moc.png 

"""

# Southern extent of Atlantic
ATLANTIC_MIN_LAT = -60
# Northern extent of Atlantic
ATLANTIC_MAX_LAT = 70
# Western extent above and below equator.
ATLANTIC_MIN_LON_ABOVE_EQ = -80
ATLANTIC_MIN_LON_BELOW_EQ = -70
# Eastern extent (includes some of Mediterranean)
ATLANTIC_MAX_LON = 20

ATLANTIC_CENTRE_LON = -30

def get_atlantic_mask(mask, lons, lats):
    """
    Return a masked array that masks the whole globe except the Atlantic.

    Input is the global bathymetry mask and the grid.
    """

    assert(len(mask.shape) == 3)

    def find_nearest_index(array, value):
        return (np.abs(array - value)).argmin()

    # Get the North and South extents. 
    north_idx = find_nearest_index(lats[:,0], ATLANTIC_MAX_LAT)
    south_idx = find_nearest_index(lats[:,0], ATLANTIC_MIN_LAT)
    equator_idx = find_nearest_index(lats[:,0], 0)
    west_a_idx = find_nearest_index(lons[0,:], ATLANTIC_MIN_LON_ABOVE_EQ)
    west_b_idx = find_nearest_index(lons[0,:], ATLANTIC_MIN_LON_BELOW_EQ)
    east_idx = find_nearest_index(lons[0,:], ATLANTIC_MAX_LON)
    centre_idx = find_nearest_index(lons[0,:], ATLANTIC_CENTRE_LON)

    # Initialise to all masked. 
    ret_mask = np.ones_like(mask, dtype=bool)

    # Start at centre_lon and go East. 
    ret_mask[:,south_idx:north_idx, centre_idx:east_idx] = \
        mask[:,south_idx:north_idx, centre_idx:east_idx]

    # Start at centre_lon and go West. This one is trickier since the
    # we treat above and below the equator differently.
    ret_mask[:,equator_idx:north_idx, west_a_idx:centre_idx] = \
        mask[:,equator_idx:north_idx, west_a_idx:centre_idx]

    ret_mask[:,south_idx:equator_idx, west_b_idx:centre_idx] = \
        mask[:,south_idx:equator_idx, west_b_idx:centre_idx]

    return ret_mask

def calc_moc(ty_trans, lons, lats):
    """
    Calculate the Global and Atlantic Meridional overturning. 
    """

    # Expect the input array to be [depth, latitude, longitude]
    assert(len(ty_trans.shape) == 3)

    # Result is 2d array depth v latitude. 
    global_moc = np.zeros_like(ty_trans[:,:,0])
    global_moc_test = np.zeros_like(global_moc)
    atlantic_moc = np.zeros_like(global_moc)

    # moc(z, lat) = integrate over longitude then integrate from
    # depth to surface for each z. 
    # Inner sum is first integral, outer is second. 
    global_moc = np.cumsum(np.sum(ty_trans, axis=2), axis=0)
    global_moc /= 1e9

    # Should be an equivalent approach. 
    # The difference is that the level ends up not being masked.
    for k in range(global_moc_test.shape[0]):
        global_moc_test[k,:] = np.sum(np.sum(ty_trans[:k+1,:,:], axis=2), axis=0)
    global_moc_test /= 1e9

    atlantic_mask = get_atlantic_mask(ty_trans.mask, lons, lats)
    # Apply atlantic mask to ty_trans
    ty_trans_atlantic = np.ma.masked_array(ty_trans, mask=atlantic_mask)
    atlantic_moc = np.cumsum(np.sum(ty_trans_atlantic, axis=2), axis=0)
    atlantic_moc /= 1e9

    return (global_moc, global_moc_test, atlantic_moc)


def plot_moc(global_moc, global_moc_test, atlantic_moc, thickness,
             lons, lats, output_file):

    depth = np.cumsum(np.mean(thickness, axis=-1), axis=0)
    y = np.amax(lats, axis=-1)
    y = np.tile(y, (depth.shape[0], 1))

    fig = plt.figure(figsize=(15, 5))
    cmap = plt.get_cmap('RdBu_r')

    fig.add_subplot(121)
    plt.title("Global MOC (Sv)")
    plt.ylabel('Depth (m)')
    plt.xlabel('Latitude')
    plt.gca().invert_yaxis()
    plt.gca().patch.set_color('0.65')
    levels = np.linspace(np.min(global_moc), np.max(global_moc), 20)
    plt.contourf(y, depth, global_moc, levels=levels, cmap=cmap)
    plt.colorbar()

    fig.add_subplot(122)
    plt.title("Atlantic MOC (Sv)")
    plt.gca().invert_yaxis()
    plt.gca().patch.set_color('0.65')
    levels = np.linspace(np.min(atlantic_moc), np.max(atlantic_moc), 20)
    plt.contourf(y, depth, atlantic_moc, levels=levels, cmap=cmap)
    plt.colorbar()

    plt.tight_layout()
    plt.savefig(output_file)

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('input_file',
                        help="The input data file in NetCDF format.")
    parser.add_argument('output_file', default='moc.png',
                        help="""The name of the output file.""")

    args = parser.parse_args()

    with nc.Dataset(args.input_file) as f:
        ty_trans = f.variables['ty_trans'][0, :]
        lons = f.variables['geolon_t'][:]
        lats = f.variables['geolat_t'][:]
        dzt = f.variables['dzt'][0, :, :, :]

    global_moc, global_moc_test, atlantic_moc = calc_moc(ty_trans, lons, lats)
    plot_moc(global_moc, global_moc_test, atlantic_moc, dzt, lons, lats,
             args.output_file)

    return 0

if __name__ == '__main__':
    sys.exit(main())
