#!/usr/bin/env python

from __future__ import print_function

import os, sys
import time
import argparse
import shutil
import netCDF4 as nc
import h5py
import numpy as np
import datetime
import multiprocessing as mp
from scipy.stats import pearsonr

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

from lib_util import make_data_cache_filename, lowess_smooth
from mom_plot_moc import calc_atlantic_moc, get_atlantic_mask, \
                            get_atlantic_moc_grid

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

def get_nh_mask(landmask, lats):
    """
    Return a masked array that masks the whole globe except Northern Hemisphere
    ocean.

    Input is the global bathymetry mask.
    """

    equator_lat_idx = find_nearest_index(lats[:, 0], 0)

    # Mask out the Southern Hemisphere
    landmask[:equator_lat_idx, :] = True
    return landmask


def get_amoc_idx_region_lat_lons():

    lat_start = 45.
    lat_end = 55.
    lon_start = -50.
    lon_end = -30.

    return lat_start, lat_end, lon_start, lon_end


def get_indices_for_amoc_idx_region(lons, lats):
    """
    Get i, j index ranges for the amoc index region (45W to 25W, 45N to 65N)
    """

    lat_start, lat_end, lon_start, lon_end = get_amoc_idx_region_lat_lons()

    lat_start_idx = find_nearest_index(lats[:, 0], lat_start)
    lat_end_idx = find_nearest_index(lats[:, 0], lat_end)
    lon_start_idx = find_nearest_index(lons[0, :], lon_start)
    lon_end_idx = find_nearest_index(lons[0, :], lon_end)

    return lat_start_idx, lat_end_idx, lon_start_idx, lon_end_idx


def max_within_region(field, min_depth, min_lat, depths, lats):
    """
    Return maximum of field deeper than min_depth and above min_lat.
    """

    lat_idx = find_nearest_index(lats[:, 0], min_lat)
    depth_idx = find_nearest_index(depths[:], min_depth)

    max = np.ma.max(field[depth_idx:, lat_idx:])
    di, li = np.unravel_index(np.argmax(field[depth_idx:, lat_idx:]), field[depth_idx:, lat_idx:].shape)
    print('AMOC maximum of {} at depth {}, lat {}'.format(max, depths[depth_idx + di], lats[lat_idx + li, 0]))

    return max


def max_at_lat(field, lat, lats):
    """
    Return maximum of field close to lat.
    """

    lat_idx = find_nearest_index(lats[:, 0], lat)

    return np.ma.max(field[:, lat_idx])


def collect_data(work_queue, results_queue, grid_def_file):

    with nc.Dataset(grid_def_file) as gf:
        areas = gf.variables['area_t'][:]

    while True:
        input_file = work_queue.get()

        # Setup some variables from the input.
        f = nc.Dataset(input_file)
        lons = f.variables['geolon_t']
        lats = f.variables['geolat_t']
        time_var = f.variables['time']
        temp_var = f.variables['temp']
        ty_trans = f.variables['ty_trans']
        depths = f.variables['st_ocean'][:]

        nh_mask = get_nh_mask(temp_var[0, 0, :, :].mask, lats)
        lat_s, lat_e, lon_s, \
            lon_e = get_indices_for_amoc_idx_region(lons, lats)

        # Set up result vars.
        t = time_var.shape[0]
        z = len(f.dimensions['st_ocean'])
        y = len(f.dimensions['yt_ocean'])
        x = len(f.dimensions['xt_ocean'])

        nh_sst_average = np.empty((t,), dtype='f') * np.nan
        nh_sst = np.ma.empty((t, y, x), dtype='f') * np.nan
        amoc_index_sst = np.empty((t,), dtype='f') * np.nan
        amoc_psi = np.ma.empty((t, z, y), dtype='f') * np.nan
        amoc_psi_max = np.empty((t,), dtype='f') * np.nan
        amoc_psi_max_at_26n = np.empty((t,), dtype='f') * np.nan

        for t in range(time_var.shape[0]):
            # Get surface temp spatial mean in the NH
            nh_sst_average[t] = np.ma.average(np.ma.masked_array(
                                                temp_var[t, 0, :, :],
                                                mask=nh_mask), weights=areas)
            # SST
            nh_sst[t, :, :] = temp_var[t, 0, :, :]

            # Spatial average SST in AMOC index region
            amoc_index_sst[t] = np.average(temp_var[t, 0, lat_s:lat_e, lon_s:lon_e],
                                        weights=areas[lat_s:lat_e, lon_s:lon_e])

            # Get AMOC psi and max within chosen region.
            amoc_psi[t, :, :] = calc_atlantic_moc(ty_trans[t, :, :, :],
                                                    lons, lats)
            amoc_psi_max[t] = max_within_region(amoc_psi[t,:,:],
                                                    500.0, 30.0, depths, lats)
            amoc_psi_max_at_26n[t] = max_at_lat(amoc_psi[t,:,:], 26.5, lats)

        f.close()
        results = (('nh_sst_average', nh_sst_average), ('nh_sst', nh_sst),
                    ('amoc_index_sst', amoc_index_sst), ('amoc_psi', amoc_psi),
                    ('amoc_psi_max', amoc_psi_max),
                    ('amoc_psi_max_at_26n', amoc_psi_max_at_26n))
        results_queue.put((input_file, results))


def get_time_var(file):

    with nc.Dataset(file) as f:
        time_var = f.variables['time'][:]

    return time_var


def preprocess_data(fname, input_files, grid_def_file):
    """
    Create a file for the preprocessed data. It will contain:
    1) timeseries of the max AMOC psi (1d),
    2) timeseries of the AMOC index (1d),
    3) timeseries of AMOC psi (3d),
    4) timeseries of SST diff to NH mean (3d)
    """

    print('master with pid {}'.format(os.getpid()))
    print('Preprocessing data ...', end='')
    start_time = time.time()

    # Get time dim for all input files
    p = mp.Pool()
    tv = p.map(get_time_var, input_files)
    p.close()
    p.join()
    tv_sizes = [a.shape[0] for a in tv]
    tv_total = sum(tv_sizes)
    assert tv_total > 0

    # Now collect data. This will be done by a pool of workers
    num_workers = max(1, mp.cpu_count() / 2)
    job_queue = mp.Queue(len(input_files))
    results_queue = mp.Queue(num_workers)

    # Start workers
    workers = []
    for w in range(num_workers):
        args = (job_queue, results_queue, grid_def_file)
        w = mp.Process(target=collect_data, args=args)
        workers.append(w)
        w.start()

    # Put all jobs in the queue
    for f in input_files:
        job_queue.put(f)

    # Setup output file
    with nc.Dataset(input_files[0]) as inf:
        z = len(inf.dimensions['st_ocean'])
        y = len(inf.dimensions['yt_ocean'])
        x = len(inf.dimensions['xt_ocean'])
        ty_trans = inf.variables['ty_trans'][0, :]
        lats = inf.variables['geolat_t'][:]
        lons = inf.variables['geolon_t'][:]

    f = h5py.File(fname)
    f.create_dataset('time', (tv_total,), dtype='f', chunks=True)
    f.create_dataset('nh_sst_average', (tv_total,), dtype='f', chunks=True)
    f.create_dataset('nh_sst', (tv_total, y, x), dtype='f', chunks=True)
    f.create_dataset('amoc_index_sst', (tv_total,), dtype='f', chunks=True)
    f.create_dataset('amoc_psi', (tv_total, z, y), dtype='f', chunks=True)
    f.create_dataset('amoc_psi_max', (tv_total,), dtype='f', chunks=True)
    f.create_dataset('amoc_psi_max_at_26n', (tv_total,), dtype='f', chunks=True)
    f.create_dataset('atlantic_mask', (z, y, x), dtype='i', chunks=True)
    f.create_dataset('amoc_psi_mask', (z, y), dtype='i', chunks=True)

    f['time'][:] = np.concatenate(tv)
    f['atlantic_mask'][:] = get_atlantic_mask(ty_trans.mask, lons, lats)
    f['amoc_psi_mask'][:] = calc_atlantic_moc(ty_trans, lons, lats).mask
    f.flush()

    # A dictionary containing the global indices for each input file
    global_idx = np.split(range(tv_total), np.cumsum(tv_sizes)[:-1])
    assert len(global_idx) == len(input_files)
    indx_dict = dict(zip(input_files, global_idx))

    # Get data from workers load into file at the correct index.
    for i in range(len(input_files)):
        input_file, result_arrays = results_queue.get()
        indx = indx_dict[input_file]

        for name, data in result_arrays:
            if len(data.shape) == 1:
                f[name][indx[0]:indx[-1]+1] = data[:]
            elif len(data.shape) == 2:
                f[name][indx[0]:indx[-1]+1, :] = data[:, :]
            elif len(data.shape) == 3:
                f[name][indx[0]:indx[-1]+1, :, :] = data[:, :, :]
            else:
                import pdb
                pdb.set_trace()
                assert False
        print('.', end='')
        sys.stdout.flush()

    # Terminate workers
    job_queue.close()
    results_queue.close()
    for w in workers:
        w.terminate()

    for w in workers:
        w.join()

    f.close()

    print(' finished in {} seconds'.format(time.time() - start_time))


def calc_sst_and_amoc_cor(args):

    # Global indices for this proc.
    my_gi_idxs, my_gj_idxs, data_file = args

    f = h5py.File(data_file, mode='r')
    nh_sst = f['/nh_sst']
    nh_sst_average = f['/nh_sst_average'][:]
    amoc_psi_max = f['/amoc_psi_max'][:]
    atlantic_mask = f['/atlantic_mask'][0, :, :]

    assert(nh_sst.shape[0] == len(amoc_psi_max))

    cor_map = np.zeros((len(my_gj_idxs), len(my_gi_idxs)))
    cor_map = np.ma.masked_array(cor_map,
                mask=atlantic_mask[np.ix_(my_gj_idxs, my_gi_idxs)])

    for j, gj in enumerate(my_gj_idxs):
        for i, gi in enumerate(my_gi_idxs):
            if not cor_map.mask[j, i]:
                cor_map[j, i] = float(pearsonr(nh_sst[:, gj, gi] - nh_sst_average,
                                                amoc_psi_max)[0])
    f.close()
    return cor_map


def sst_and_amoc_max_cor_map(data_file, use_smoothed=False):

    with h5py.File(data_file, 'r') as f:
        nh_sst = f['/nh_sst']
        i_idxs = range(nh_sst.shape[2])
        j_idxs = range(nh_sst.shape[1])

    # Split the j_idxs over available procs.
    j_chunks = np.array_split(j_idxs, 20)

    import time
    start_time = time.time()

    pool = mp.Pool()
    pargs = zip([i_idxs]*len(j_chunks), j_chunks, [data_file]*len(j_chunks))
    results = pool.map(calc_sst_and_amoc_cor, pargs)
    cor_map = np.ma.vstack(results)

    print('sst_and_amoc_max_cor_map took {} seconds'.format(time.time() - start_time))

    return cor_map


def calc_amoc_psi_and_index_cor(args):
    """
    A depth plot of correlation with the AMOC index.
    """

    # Global indices for this proc.
    my_gy_idxs, my_gz_idxs, data_file = args

    f = h5py.File(data_file, mode='r')
    amoc_psi = f['/amoc_psi']
    amoc_index_sst = f['/amoc_index_sst'][:]
    nh_sst_average = f['/nh_sst_average'][:]
    amoc_psi_mask = f['/amoc_psi_mask'][:, :]

    amoc_idx = amoc_index_sst - nh_sst_average

    assert(amoc_psi.shape[0] == len(amoc_index_sst))

    cor_map = np.zeros((len(my_gz_idxs), len(my_gy_idxs)))
    cor_map = np.ma.masked_array(cor_map,
                mask=amoc_psi_mask[np.ix_(my_gz_idxs, my_gy_idxs)])

    for z, gz in enumerate(my_gz_idxs):
        for y, gy in enumerate(my_gy_idxs):
            if not cor_map.mask[z, y]:
                cor_map[z, y] = float(pearsonr(amoc_psi[:, gz, gy],
                                                amoc_index_sst)[0])
    f.close()
    return cor_map


def amoc_psi_and_amoc_index_cor_map(data_file, use_smoothed=False):

    with h5py.File(data_file, 'r') as f:
        amoc_psi = f['/amoc_psi']
        y_idxs = range(amoc_psi.shape[2])
        z_idxs = range(amoc_psi.shape[1])

    # Split the z_idxs over available procs.
    z_chunks = np.array_split(z_idxs, 20)

    import time
    start_time = time.time()

    pool = mp.Pool()
    pargs = zip([y_idxs]*len(z_chunks), z_chunks, [data_file]*len(z_chunks))
    results = pool.map(calc_amoc_psi_and_index_cor, pargs)
    cor_map = np.ma.vstack(results)

    print('calc_amoc_psi_and_index_cor took {} seconds'.format(time.time() - start_time))

    return cor_map


def plot_timeseries(data_file, output_file, offset=0):

    # Read data and plot
    with h5py.File(data_file) as f:
        time = f['/time'][:]
        amoc_index_sst = f['/amoc_index_sst'][:]
        nh_sst_average = f['/nh_sst_average'][:]
        amoc_psi_max = f['/amoc_psi_max'][:]
        amoc_psi_max_at_26n = f['/amoc_psi_max_at_26n'][:]

    start_date = datetime.datetime(year=1, month=1, day=1)
    time_dim = [start_date + datetime.timedelta(days=int(t)) for t in time]
    amoc_index = amoc_index_sst - nh_sst_average

    if offset != 0:
        time_dim = time_dim[:-offset]
        amoc_index = amoc_index[:-offset]
        amoc_psi_max = amoc_psi_max[offset:]

    # Do some smoothing
    amoc_index_sm = lowess_smooth(amoc_index)
    amoc_psi_max_anomaly = amoc_psi_max - np.mean(amoc_psi_max)
    amoc_psi_max_anomaly_sm = lowess_smooth(amoc_psi_max_anomaly)
    amoc_psi_max_sm = lowess_smooth(amoc_psi_max)

    # Time series of the maximum overturning stream function anomaly
    fig, axarr = plt.subplots(2, sharex=True)
    axarr[0].plot(time_dim, amoc_psi_max_anomaly, 'red', lw=1.0, alpha=0.5)
    axarr[0].plot(time_dim, amoc_psi_max_anomaly_sm, 'darkred', lw=3.0)
    axarr[0].set_ylabel('AMOC Maximum Anomaly (Sv)', color='darkred')
    axarr[0].set_title('Connection between Max Overturning Streamfunction and AMOC Index')

    # Time series of the AMOC index
    ax2 = axarr[0].twinx()
    ax2.plot(time_dim, amoc_index, 'blue', lw=1.0, alpha=0.5)
    ax2.plot(time_dim, amoc_index_sm, 'darkblue', lw=3.0)
    ax2.set_ylabel('AMOC Index (K)', color='darkblue')

    p_corr_sm = pearsonr(amoc_psi_max_anomaly_sm, amoc_index_sm)
    print('p_corr_sm: {}'.format(p_corr_sm))
    p_corr = pearsonr(amoc_psi_max_anomaly, amoc_index)
    print('p_corr: {}'.format(p_corr))
    text = 'Pearson correlation (smoothed): {0:.2f} ({1:.2f})'. \
                format(p_corr[0], p_corr_sm[0])
    print(text)

    props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
    axarr[0].text(0.05,  0.95, text, transform=axarr[0].transAxes,
                    verticalalignment='top', bbox=props)

    # Time series of AMOC mean
    time_dim = [start_date + datetime.timedelta(days=int(t)) for t in time]
    amoc_psi_max_at_26n_sm = lowess_smooth(amoc_psi_max_at_26n)
    axarr[1].plot(time_dim, amoc_psi_max_at_26n, 'green', lw=1.0)
    axarr[1].plot(time_dim, amoc_psi_max_at_26n_sm, 'darkgreen', lw=3.0)
    axarr[1].set_ylabel('Sv', color='darkgreen')
    axarr[1].set_xlabel('Year')
    axarr[1].set_title('AMOC monthly maximum at 26.5N')

    fig.set_size_inches(14,10)
    plt.tight_layout()
    plt.savefig(output_file)
    plt.close()


def plot_surface_cor_map(cor_map, an_input_file, output_file):

    with nc.Dataset(an_input_file) as f:
        lats = f.variables['geolat_t'][:]
        lons = f.variables['geolon_t'][:]

    # Mask out lats, lons with the same mask as used by cor_map
    mlats = np.ma.masked_array(lats, mask=cor_map.mask)
    mlons = np.ma.masked_array(lons, mask=cor_map.mask)

    m = Basemap(projection='mill', lat_ts=10, llcrnrlon=np.min(mlons),
                urcrnrlon=np.max(mlons), llcrnrlat=np.min(mlats),
                urcrnrlat=np.max(mlats), resolution='c')
    x, y = m(lons, lats)

    cs = m.pcolormesh(x, y, cor_map, cmap=plt.cm.coolwarm)
    m.fillcontinents()
    m.drawmapboundary()
    plt.colorbar()
    plt.title('Correlating max AMOC streamfunction\n and SST difference from NH mean')
    plt.savefig(output_file)
    plt.close()


def plot_depth_cor_map(cor_map, an_input_file, output_file):

    with nc.Dataset(an_input_file) as f:
        lats = f.variables['geolat_t'][:]
        lons = f.variables['geolon_t'][:]
        dzt = f.variables['dzt']

        depth, north_idx, south_idx, _, _ = \
            get_atlantic_moc_grid(dzt[0, :, :, :], lons, lats)

    new_lats = lats[south_idx:north_idx, :]
    depth = depth[:, south_idx:north_idx]
    cor_map = cor_map[:, south_idx:north_idx]

    y = np.amax(new_lats,  axis=-1)
    y = np.tile(y, (depth.shape[0], 1))

    fig = plt.figure(figsize=(10, 5))
    plt.title("Correlation between Overturning Streamfunction\n and AMOC Index")
    plt.ylabel('Depth (m)')
    plt.xlabel('Latitude')
    plt.gca().invert_yaxis()
    plt.gca().patch.set_color('0.65')
    levels = np.linspace(np.min(cor_map), np.max(cor_map), 20)
    cmap = plt.get_cmap('coolwarm')
    plt.contourf(y, depth, cor_map, levels=levels, cmap=cmap)
    plt.colorbar()
    plt.savefig(output_file)
    plt.close()


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('input_files', nargs='+',
                        help='The MOM ocean.nc input data files.')
    parser.add_argument('--grid_def_file', default='ocean_grid.nc',
                        help='The ocean grid definition file.')
    parser.add_argument('--output_file', default='amoc_index_timeseries.png',
                        help='Name of the plot.')
    plot_name = 'amoc_psi_and_sst_correlation_surface_map.png'
    parser.add_argument('--surface_correlation_output_file', default=plot_name,
                        help='Name of the plot.')
    plot_name = 'amoc_psi_and_index_correlation_depth_map.png'
    parser.add_argument('--depth_correlation_output_file', default=plot_name,
                        help='Name of the plot.')
    parser.add_argument('--offset', default=0, type=int, help='FIXME.')

    args = parser.parse_args()

    assert(len(args.input_files) == len(set(args.input_files)))

    do_timeseries_plot = args.output_file.lower() != 'none'
    do_surface_correlation_plot = \
        args.surface_correlation_output_file.lower() != 'none'
    do_depth_correlation_plot = \
        args.depth_correlation_output_file.lower() != 'none'

    data_file = make_data_cache_filename(__file__, main,
                                    args.input_files, suffix='hdf5')

    if not os.path.exists(data_file):
        preprocess_data(data_file, args.input_files, args.grid_def_file)

    if do_timeseries_plot:
        plot_timeseries(data_file, args.output_file, offset=args.offset)

    if do_surface_correlation_plot:
        surface_cor_map = sst_and_amoc_max_cor_map(data_file, use_smoothed=True)
        plot_surface_cor_map(surface_cor_map, args.input_files[0],
                        args.surface_correlation_output_file)

    if do_depth_correlation_plot:
        depth_cor_map = amoc_psi_and_amoc_index_cor_map(data_file,
                                                    use_smoothed=True)
        plot_depth_cor_map(depth_cor_map, args.input_files[0],
                            args.depth_correlation_output_file)

    return 0

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