"""

Analysing Scales of Precipitation diagnostics
Space-time Coherence package

These functions allow the user to compute and plot diagnostics of the
spatial and temporal coherence of precipitation in observations and/or
models.  These functions should work for any resolution or timescale.

These functions assume that the input precipitation dataset is in netCDF.

Written by Nicholas Klingaman
nicholas.klingaman@ncas.ac.uk

(C) The author 2017

"""
import csv
import os
import json
import numpy as np
import iris
import cf,cfplot as cfp
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
from six.moves import range

def update_output_json(json_key, data_description, wk_dir):
    """Add the dictionary 'data_description' under the key 'json_key' in
    the descriptive output json"""
    json_filename = os.path.join(wk_dir,"output.json")
    # Only update output json if it exists
    if os.path.exists(json_filename):
        with open(json_filename,"r") as output_json:
            output=json.load(output_json)
        output[json_key].update(data_description)
        with open(json_filename,"w") as output_json:
            json.dump(output,output_json, indent=2)

def default_options():

    """
    Dummy function to return initial values of key parameters that must be defined.
    """

    region_size=0
    input_box_size=0
    lag_length=0
    autocorr_length=0
    time_type='None'
    grid_type='None'
    time_desc='None'
    grid_desc='None'
    model_desc='None'
    return(region_size,input_box_size,lag_length,autocorr_length,time_type,grid_type,time_desc,grid_desc,model_desc)

def parameters():

    """
    Parameters that control the size of arrays used in other functions.
    """

    max_box_distance=100
    max_timesteps=100
    max_boxes=100
    return(max_box_distance,max_boxes,max_timesteps)

def read_precip(model_dict):

    """
    Use iris to read precipitation data into a Cube.
    The area to read is controlled by the "region" element of the dataset dictionary (model_dict)
    "region" should be defined as min_lat, max_lat, min_lon, max_lon.

    Arguments:
    * model_dict: The dictionary containing details of the dataset to read.
      (See the create_dict function).

    Returns:
    * precip: An iris cube of precipitation data.
    """

    constraint = iris.Constraint(model_dict['constraint'],
                                     latitude=lambda cell: model_dict['region'][0] <= cell <= model_dict['region'][1],
                                     longitude=lambda cell: model_dict['region'][2] <= cell <= model_dict['region'][3])
    precip = iris.load_cube(model_dict['infile'],constraint)*model_dict['scale_factor']
    try:
        precip.coord('latitude').guess_bounds()
    except:
        pass
    try:
        precip.coord('longitude').guess_bounds()
    except:
        pass
    return(precip)

def compute_histogram(precip,bins):

    """
    Computes 1D and 2D histograms of precipitation from an input iris cube of precipitation.
    The 2D histogram is the histogram of precipitation on consecutive timesteps.

    Arguments:
    * precip :
       An iris cube of precipitation
    * bins:
       A numpy array of the edges of the bins for which to compute the histogram,
       in the same units as "precip" above.

    Returns:
    * oned_hist:
       The one-dimensional histogram, as a numpy array, normalised so that the sum of the values is one.
    * twod_hist:
       The two-dimensional histogram, as a numpy array, normalised so that the sum of the values is one.
    """

    oned_hist, bin_edges = np.histogram(precip.data,bins)
    nbins=len(oned_hist)
    twod_hist=np.zeros((nbins,nbins))
    print('---> Computing 2D histogram')
    for t, t_slice in enumerate(precip.slices(['time'])):
        next_slice = t_slice.copy()
        next_slice.data = np.roll(t_slice.data,1,0)
        twod_hist_temp,xedges,yedges = np.histogram2d(t_slice.data,next_slice.data,bins)
        twod_hist=twod_hist+twod_hist_temp
    twod_hist=twod_hist/np.sum(twod_hist)
    oned_hist=oned_hist/np.float(np.sum(oned_hist))
    return(oned_hist,twod_hist)

def plot_histogram(oned_hist,twod_hist,model_dict,bins,title=True,colorbar=True,wk_dir='.',ext='.ps'):

    """
    Creates a PostScript plot of the 1D and 2D histograms calculated in compute_histogram.

    Arguments:
    * oned_hist:
       The 1D histogram computed in compute_histogram.
    * twod_hist:
       The 2D histogram computed in compute_histogram.
    * model_dict:
       The dictioary containing details of this model.
    * bins:
       The edges of the histogram bins, as a numpy array.

    Optional arguments:
    * title:
       A logical to control whether the title is printed at the top of the plot (useful for creating multi-panel plots).
    * colorbar:
       A logical to control whether the colorbar is printed at the bottom of the plot (useful for creating multi-panel plots).
    """

    print('---> Plotting 2D histogram')
    nbins = np.size(oned_hist)
    hist_con_levs=[1e-5,2e-5,4e-5,7e-5,1e-4,2e-4,4e-4,7e-4,1e-3,2e-3,4e-3,7e-3,1e-2,2e-2,4e-2,7e-2,1e-1]
    fig = plt.figure()
    ax = fig.add_subplot(111)
    cmap = plt.cm.get_cmap("viridis_r")
    norm = BoundaryNorm(hist_con_levs,ncolors=cmap.N,clip=True)
    contour = ax.pcolormesh(np.arange(nbins+1),np.arange(nbins+1),twod_hist,cmap=cmap,norm=norm)
    if colorbar == True:
        cbar = fig.colorbar(contour,orientation='horizontal',ticks=hist_con_levs)
        cbar.ax.set_xlabel('Probability',fontsize=18)
        cbar.ax.set_xticklabels(['1e-5','2e-5','4e-5','7e-5','1e-4','2e-4','4e-4','7e-4','1e-3','2e-3','4e-3','7e-3','1e-2','2e-2','4e-2','7e-2','1e-1'])
    ax.set_xlabel('Precipitation at time t (mm day$^{-1}$)',fontsize=16)
    ax.set_ylabel('Precipitation at time t+1 (mm day$^{-1}$)',fontsize=16)
    ticklabels=['< '+str(bins[1])]
    for bin in range(1,nbins):
        ticklabels.append(str(bins[bin]))
    ticklabels.append(' > '+str(bins[nbins-1]))
    ax.set_xticks(np.arange(nbins+1))
    ax.set_xticklabels(ticklabels,fontsize=14)
    ax.set_yticks(np.arange(nbins+1))
    ax.set_yticklabels(ticklabels,fontsize=14)
    if title == True:
        title_string = '2D histogram for '+model_dict['legend_name']
        if 'time_desc' in model_dict:
            title_string = title_string+' '+model_dict['time_desc']
        if 'grid_desc' in model_dict:
            title_string = title_string+' '+model_dict['grid_desc']
        if 'region_desc' in model_dict:
            title_string = title_string+' - '+model_dict['region_desc']
        title_string = title_string + ' data'
        ax.set_title(title_string)
    ax.axis([0,nbins,0,nbins])
    ax.set_xlim(xmin=0,xmax=nbins)
    ax2 = ax.twinx()
    ax2.plot(np.arange(nbins)+0.5,oned_hist,'k--',marker='o',markersize=8)
    ax2.set_yscale('log',nonposy='clip')
    ax2.set_ylim(ymin=0.0009,ymax=1.0)
    ax2.set_ylabel('Probability of precipitation in bin',fontsize=18)
    ax2.set_yticks([1e-3,1.4e-3,2e-3,3e-3,4.5e-3,7e-3,1e-2,1.4e-2,2e-2,3e-2,4.5e-2,7e-2,1e-1,1.4e-1,2e-1,3e-1,4.5e-1,7e-1,1])
    ax2.set_yticklabels(['1.0e-3','1.4e-3','2.0e-3','3.0e-3','4.5e-3','7.0e-3','1.0e-2','1.4e-2','2.0e-2','3.0e-2','4.5e-2',
                         '7.0e-2','1.0e-1','1.4e-1','2.0e-1','3.0e-1','4.5e-1','7.0e-1','1.0e0'],fontsize=14)
    ax2.set_xlim(xmin=0,xmax=nbins)
    plot_name=wk_dir+'/asop_coherence.'+model_dict['name']
    if 'grid_type' in model_dict:
        plot_name=plot_name+'_'+model_dict['grid_type']
    if 'time_type' in model_dict:
        plot_name=plot_name+'_'+model_dict['time_type']
    if 'region_name' in model_dict:
        plot_name=plot_name+'_'+model_dict['region_name'].replace(" ","_")
    plot_name=plot_name+'_precip_twodpdf'+ext
    plt.savefig(plot_name,bbox_inches='tight')

    desc = {os.path.relpath(plot_name,wk_dir): {
        "longname": os.path.basename(plot_name).split(".")[1].replace("_"," "),
        "description": "1-d and 2-d histogram plots"
        }
    }
    update_output_json("plots",desc,wk_dir)

def haversine(origin, destination):

    """
    Compute distance between two gridpoints in km.

    Method: Haversine function

    Arguments:
    * origin: Tuple of (latitude,longitude) for origin point
    * destination: Tuple of (latitude, longitude) for destination point

    Returns:
    * d: distance between origin and destination in km
    """

    import math

    lat1, lon1 = origin
    lat2, lon2 = destination
    radius = 6371 # km

    dlat = math.radians(lat2-lat1)
    dlon = math.radians(lon2-lon1)
    a = math.sin(dlat/2) * math.sin(dlat/2) + math.cos(math.radians(lat1)) \
        * math.cos(math.radians(lat2)) * math.sin(dlon/2) * math.sin(dlon/2)
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1-a))
    d = radius * c
    return d

def compute_equalgrid_corr(precip,model_dict,metrics_csv=None,wk_dir='.'):

    """
    Compute correlations in space and time, using the native spatial and temporal
    resolutions of the input data.

    Method:
    The input spatial domain is broken down into non-overlapping regions of N x N gridpoints, where
    N is controlled by the "region_size" input parameter.

    Correlations in both space and time are computed with respect to lag-0 at the central
    point in each region.

    Composite (average) correlations are computed over all regions in the domain.

    To express the average correlations as a function of distance from the central point, bins of
    distance from the central point are created, with width delta_x (the x spacing of the input data,
    taken from the dataset dictionary), starting from 0.5*delta_x.  Correlations are averaged within each bin,
    at each lag.

    Arguments:
     * precip:
        An iris cube of precipitation data
     * model_dict:
        The dictionary containing information about this dataset.

    Returns:
     * corr_map (lag_length,region_size,region_size):
        A 'map' of the composite correlations over all regions in the domain, at each lag.
     * lag_vs_distance (lag_length,region_size):
        The composite correlation over all regions in the domain, expressed a function of time (lag)
        and distance from the central point, in bins of delta_x (x spacing of the input data) starting
        from 0.5*delta_x.
     * autocorr (lag_length):
        The composite auto-correlation of precipitation at the central point, averaged over all regions
        in the domain.
     * npts_map (lag_length,region_size,region_size):
        The number of gridpoints that contributed to the composite map of correlations.  Can be used
        to weight the correlation map if averaging over multiple maps.
     * npts (lag_length,region_size):
        The number of gridpoints in each lag and distance bin of the lag_vs_distance array.  Used by
        the corresponding plotting function to determine whether there are any points in this bin.
    """

    lag_length = model_dict['lag_length']
    region_size = model_dict['region_size']
    print('---> Computing correlations for '+str(region_size)+'x'+str(region_size)+' sub-regions')
    latitude = precip.coord('latitude').points
    longitude = precip.coord('longitude').points
    nlon=len(longitude)
    nlat=len(latitude)
    print('----> Info: Size of domain in native gridpoints: '+str(nlon)+' longitude x '+str(nlat)+' latitude.')
    nregions=0
    npts=np.zeros((lag_length,region_size),dtype=np.int32)
    npts_map=np.zeros((lag_length,region_size,region_size),dtype=np.int32)
#    distance_x=np.zeros(region_size)
#    distance_y=np.zeros(region_size)
    corr_map=np.zeros((lag_length,region_size,region_size))
    lag_vs_distance=np.zeros((lag_length,region_size))
    autocorr=np.zeros((lag_length))
    for region_xstart in range(0,nlon-region_size+1,region_size):
        region_xcentre = region_xstart+region_size//2
        for region_ystart in range(0,nlat-region_size+1,region_size):
            region_ycentre = region_ystart+region_size//2
            central_precip = precip[:,region_ycentre,region_xcentre]
#            for region_x in range(region_size):
#                distance_x[region_x]=np.abs(region_xstart+region_x-region_xcentre)
#            for region_y in range(region_size):
#                distance_y[region_y]=np.abs(region_ystart+region_y-region_ycentre)*model_dict['dy']/float(model_dict['dx'])
            for region_x in range(region_size):
                for region_y in range(region_size):
#                    distance=np.int(round(np.sqrt(distance_x[region_x]*distance_x[region_x]+distance_y[region_y]*distance_y[region_y])))-1
                    km_distance = haversine((latitude[region_ycentre],longitude[region_xcentre]),(latitude[region_ystart+region_y],longitude[region_xstart+region_x]))
                    distance = np.int(round(km_distance/model_dict['dx']))-1
                    if distance < region_size:
                        remote_precip=precip[:,region_y+region_ystart,region_x+region_xstart]
                        corr = np.corrcoef([central_precip.data,remote_precip.data])[1,0]
                        if not np.isnan(corr):
                            corr_map[0,region_y,region_x]=corr+corr_map[0,region_y,region_x]
                            npts_map[0,region_y,region_x]=npts_map[0,region_y,region_x]+1
                        if (region_x + region_xstart == region_xcentre) and (region_y + region_ystart == region_ycentre):
                            autocorr[0]=autocorr[0]+1
                            for lag in range(1,lag_length):
                                corr = np.corrcoef(central_precip.data,np.roll(central_precip.data,lag,0))[1,0]
                                if not np.isnan(corr):
                                    corr_map[lag,region_y,region_x]=corr+corr_map[lag,region_y,region_x]
                                    npts_map[lag,region_y,region_x]=npts_map[lag,region_y,region_x]+1
                                    autocorr[lag]=corr+autocorr[lag]
                                    npts[lag,0]=npts[lag,0]+1
                        else:
                            if not np.isnan(corr):
                                lag_vs_distance[0,distance]=lag_vs_distance[0,distance]+corr
                                npts[0,distance]=npts[0,distance]+1
                            for lag in range(1,lag_length):
                                corr = np.corrcoef([central_precip.data,np.roll(remote_precip.data,lag,0)])[1,0]
                                if not np.isnan(corr):
                                    corr_map[lag,region_y,region_x] = corr+corr_map[lag,region_y,region_x]
                                    npts_map[lag,region_y,region_x] = npts_map[lag,region_y,region_x]+1
                                    lag_vs_distance[lag,distance] = corr+lag_vs_distance[lag,distance]
                                    npts[lag,distance]=npts[lag,distance]+1
            nregions = nregions+1
    corr_map = corr_map/npts_map
    print('----> Info: There are '+str(nregions)+' '+str(region_size)+'x'+str(region_size)+' sub-regions in your input data.')
    for lag in range(lag_length):
        for dist in range(region_size):
            if npts[lag,dist] > 0:
                lag_vs_distance[lag,dist]=lag_vs_distance[lag,dist]/npts[lag,dist]
            # If there are no gridpoints in range, set correlation to a missing value
            if npts[lag,dist] == 0:
                lag_vs_distance[lag,dist]=-999
        autocorr[lag]=autocorr[lag]/nregions

    # Write metrics to file for CMEC driver
    if metrics_csv is not None:
        box_length_x = np.int(model_dict['box_size']//model_dict['dx'])
        box_length_y = np.int(model_dict['box_size']//model_dict['dy'])
        data_to_csv = {"Dataset": model_dict["name"],
            "dx": model_dict["dx"],
            "dy": model_dict["dy"],
            "region_size": region_size,
            "nlon": nlon,
            "nlat": nlat,
            "nregions": nregions,
            "gridbox_x": model_dict["box_size"],
            "gridbox_y": model_dict["box_size"],
            "nx": box_length_x,
            "ny": box_length_y}
        field_names = ["Dataset",
            "dx",
            "dy",
            "region_size",
            "nlon",
            "nlat",
            "nregions",
            "gridbox_x",
            "gridbox_y",
            "nx",
            "ny"]
        if not os.path.exists(metrics_csv):
            with open(metrics_csv,"w") as fname:
                csv_writer = csv.DictWriter(fname, fieldnames=field_names)
                csv_writer.writeheader()
                csv_writer.writerow(data_to_csv)
        else:
            with open(metrics_csv,"a") as fname:
                csv_writer = csv.DictWriter(fname, fieldnames=field_names)
                csv_writer.writerow(data_to_csv)

        data_desc = {
            os.path.basename(metrics_csv): {
                "longname": os.path.basename(metrics_csv).split(".")[1].replace("_"," "),
                "description": "Box dimension information organized by model"
            }
        }
        update_output_json("metrics",data_desc,wk_dir)

    return(corr_map,lag_vs_distance,autocorr,npts_map,npts)

def plot_equalgrid_corr(corr_map,lag_vs_distance,autocorr,npts,model_dict,title=True,colorbar=True,wk_dir='.',ext='.ps'):

    """
    Plots correlations as functions of space and time, which were first computed
    using compute_equalgrid_corr.

    Two types of plots are created:

        1. For each lag (from 0 to lag_length defined in model_dict), a 2D map
        of the composite correlations against the central point (at lag 0) for all
        points in the region (of length region_size).

        2. A single lag vs. distance plot showing the composite correlations against
        the central point (at lag 0), averaged over all points in each region in each
        distance bin (in steps of dx starting at 0.5dx), as well as the auto-correlation
        at the central point.

    See Fig. 2 in Klingaman et al. (2017, GMD, doi:10.5194/gmd-10-57-2017) for examples
    of these diagrams.

    Arguments:
      * corr_map (lag_length,region_size,region_size):
         Composite maps of correlations at each lag, returned from compute_equalgrid_corr
      * lag_vs_distance (lag_length,region_size):
         Composite correlations over all regions in the domain, expressed as a function of
         time (lag) and distance from the central point, returned from compute_equalgrid_corr
      * autocorr (lag_length):
         The composite auto-correlation of precipitation at the central point, averaged over
         all regions in the domain.
      * npts (lag_length,region_size):
         The number of gridpoints in each distance bin of the lag_vs_distance array.  Used to
         to determine whether there are any points in each distance bin.

    Optional arguments:
      * title:
         Include a title on the plot
      * colorbar:
         Include a colorbar on the plot

    Returns:
      None
    """

    region_size = model_dict['region_size']
    lag_length = model_dict['lag_length']

    print('---> Plotting correlation maps for '+str(region_size)+'x'+str(region_size)+' sub-regions')
    corr_con_levs=[0.05,0.15,0.25,0.35,0.45,0.55,0.65,0.75,0.85,0.95]

    # Plot correlation maps at each lag
    for lag in range(lag_length):
        plot_name=wk_dir+'/asop_coherence.'+model_dict['name']
        if 'grid_type' in model_dict:
            plot_name=plot_name+'_'+model_dict['grid_type']
        if 'time_type' in model_dict:
            plot_name=plot_name+'_'+model_dict['time_type']
        if 'region_name' in model_dict:
            plot_name=plot_name+'_'+model_dict['region_name'].replace(" ","_")
        plot_name=plot_name+'_precip_'+str(region_size)+'x'+str(region_size)+'maps_lag'+str(lag)+ext
        cfp.setvars(file=plot_name,text_fontsize=18,axis_label_fontsize=18,colorbar_fontsize=18)
        cfp.gopen(figsize=[8,8])
        cfp.gset(xmin=0,xmax=region_size,ymin=0,ymax=region_size)
        cfp.levs(manual=np.array(corr_con_levs))
        cfp.cscale(cmap='parula',reverse=1,ncols=len(corr_con_levs)+1,white=0)
        #cfp.axes(xticks=np.arange(region_size)+0.5,yticks=np.arange(region_size)+0.5,
        #         xticklabels=np.arange(region_size)-region_size//2,yticklabels=np.arange(region_size)-region_size//2)
        if title == True:
            title_string = 'Correlation map for '+model_dict['legend_name']
            if 'time_desc' in model_dict:
                title_string = title_string + ' - ' +model_dict['time_desc']
            if 'grid_desc' in model_dict:
                title_string = title_string + ' - ' +model_dict['grid_desc']
            if 'region_desc' in model_dict:
                title_string = title_string+' - '+model_dict['region_desc']
            title_string = title_string + ' - Lag '+str(lag)
        else:
            title_string = ''
        if colorbar == True:
            cfp.con(f=corr_map[lag,:,:],x=np.arange(region_size)+0.5,y=np.arange(region_size)+0.5,blockfill=1,
                    lines=False,line_labels=False,ptype=0,colorbar=1,colorbar_title='Correlation with (0,0) at lag 0, mean over all sub-regions',
                    title=title_string, colorbar_text_up_down=True,
                    xlabel=r'$\Delta$x (approximately '+str(round(model_dict['dx']))+' km)',
                    ylabel=r'$\Delta$y (approximately '+str(round(model_dict['dy']))+' km)',
                    xticks=np.arange(region_size)+0.5,yticks=np.arange(region_size)+0.5,
                    xticklabels=np.arange(region_size)-region_size//2,yticklabels=np.arange(region_size)-region_size//2)
        else:
            cfp.con(f=corr_map[lag,:,:],x=np.arange(region_size)+0.5,y=np.arange(region_size)+0.5,blockfill=1,
                    lines=False,line_labels=False,ptype=0,colorbar=0,title=title_string,
                    xlabel=r'$\Delta$x (approximately '+str(model_dict['dx'])+' km)',
                    ylabel=r'$\Delta$y (approximately '+str(model_dict['dy'])+' km)',
                    xticks=np.arange(region_size)+0.5,yticks=np.arange(region_size)+0.5,
                    xticklabels=np.arange(region_size)-region_size//2,yticklabels=np.arange(region_size)-region_size//2)
        for region_x in range(region_size):
            for region_y in range(region_size):
                if corr_map[lag,region_y,region_x] > 0.5:
                    cfp.plotvars.plot.text(region_x+0.5,region_y+0.5,str(corr_map[lag,region_y,region_x])[0:4],
                                           horizontalalignment='center',color='white',fontsize=20)
                elif corr_map[lag,region_y,region_x] < 0.0:
                    cfp.plotvars.plot.text(region_x+0.5,region_y+0.5,str(corr_map[lag,region_y,region_x])[0:5],
                                           horizontalalignment='center',color='black',fontsize=20)
                else:
                    cfp.plotvars.plot.text(region_x+0.5,region_y+0.5,str(corr_map[lag,region_y,region_x])[0:4],
                                           horizontalalignment='center',color='black',fontsize=20)
        cfp.gclose()

        desc = {os.path.relpath(plot_name,wk_dir): {
            "longname": os.path.basename(plot_name).split(".")[1].replace("_"," "),
            "description": "2D map of the composite correlations against the central point (at lag 0)"+
                " for all points in the region for lag {0}".format(str(lag))
            }
        }
        update_output_json("plots",desc,wk_dir)

    # Plot correlation vs. distance diagram
    print('---> Plotting lag vs. distance diagram')
    plot_name=wk_dir+'/asop_coherence.'+model_dict['name']
    if 'grid_type' in model_dict:
        plot_name=plot_name+'_'+model_dict['grid_type']
    if 'time_type' in model_dict:
        plot_name=plot_name+'_'+model_dict['time_type']
    if 'region_name' in model_dict:
        plot_name=plot_name+'_'+model_dict['region_name'].replace(" ","_")
    plot_name=plot_name+'_precip_'+str(region_size)+'x'+str(region_size)+'_lag'+str(lag_length)+ext
    cfp.setvars(file=plot_name,text_fontsize=20,axis_label_fontsize=18)
    cfp.gopen(figsize=[9,8])
    ticklabels=['Centre','0.5']
    max_dist=0
    for dist in range(2,region_size):
        if npts[0,dist] > 0 :
            ticklabels.append(str(dist-0.5))
            max_dist=dist
    ticklabels.append(str(max_dist+0.5))
    ticklabels.append(str(max_dist+1.5))
    lag_vs_distance=np.insert(lag_vs_distance,0,autocorr,1)
    cfp.gset(xmin=0,xmax=max_dist+1,ymin=0,ymax=lag_length)
    cfp.levs(manual=np.array(corr_con_levs))
    cfp.cscale(cmap='parula',reverse=1,ncols=len(corr_con_levs)+1,white=0)
    xtickvals=np.arange(max_dist+1)+2.0
    xtickvals=np.insert(xtickvals,0,[0.5,1.0])
    cfp.axes(xticks=xtickvals,yticks=np.arange(lag_length)+0.5,xticklabels=ticklabels,yticklabels=np.arange(lag_length),
             xlabel=r'$\Delta$x bins ($\Delta$x approximately '+str(round(model_dict['dx']))+' km at the equator)',
             ylabel='Lag')

    if title == True:
        title_string = 'Correlation map for '+model_dict['legend_name']
        if 'time_desc' in model_dict:
            title_string = title_string + ' - ' +model_dict['time_desc']
        if 'grid_desc' in model_dict:
            title_string = title_string + ' - ' +model_dict['grid_desc']
        if 'region_name' in model_dict:
            title_string = title_string+' - '+model_dict['region_name']
    else:
        title_string = ''
    if colorbar == True:
        cfp.con(f=lag_vs_distance[:,0:max_dist+2],x=np.arange(max_dist+2)+0.5,y=np.arange(lag_length)+0.5,blockfill=1,
                lines=False,line_labels=False,ptype=0,colorbar_title='Correlation with centre at lag=0, mean over all sub-regions',
                title=title_string)

    else:
        cfp.con(f=lag_vs_distance[:,0:max_dist+2],x=np.arange(max_dist+2)+0.5,y=np.arange(lag_length)+0.5,blockfill=1,
                lines=False,line_labels=False,ptype=0,colorbar=0,title=title_string)

    for dist in range(max_dist+2):
        for lag in range(lag_length):
            if lag_vs_distance[lag,dist] == -999:
                print('-999')
            #                cfp.plotvars.plot.text(dist+0.5,lag+0.5,'XXX',horizontalalignment='center',color='black',fontsize=20,verticalalignment='center')
            elif lag_vs_distance[lag,dist] > 0.5:
                cfp.plotvars.plot.text(dist+0.5,lag+0.5,str(lag_vs_distance[lag,dist])[0:4],
                                       horizontalalignment='center',color='white',fontsize=20)
            elif lag_vs_distance[lag,dist] < 0.0:
                cfp.plotvars.plot.text(dist+0.5,lag+0.5,str(lag_vs_distance[lag,dist])[0:5],
                                       horizontalalignment='center',color='black',fontsize=20)
            else:
                cfp.plotvars.plot.text(dist+0.5,lag+0.5,str(lag_vs_distance[lag,dist])[0:4],
                                       horizontalalignment='center',color='black',fontsize=20)
    cfp.gclose()

    desc = {os.path.relpath(plot_name,wk_dir): {
        "longname": os.path.basename(plot_name).split(".")[1].replace("_"," "),
        "description": "Correlation vs distance diagram for lag "+str(lag)
        }
    }
    update_output_json("plots",desc,wk_dir)

def compute_equalarea_corr(precip,model_dict):

    """
    Computes spatial correlations as a function of physical distance (in km).  Note that unlike
    compute_equalgrid_corr, this routine does *not* compute lagged correlations; it computes only
    instantaneous correlations.

    Method:
    As in compute_equalgrid_corr, the analysis domain is divided into square sub-regions that are
    box_size in length.  Precipitation at each point in the sub-region is correlated against the
    central point in the domain.  These correlations are binned by the physical distance from the
    central point, using bins delta_x wide starting from 0.5*delta_x.  Correlations are averaged
    within the bin and across all sub-regions.

    As of vn1.0.2, distances are computed by a Haversine function rather than by using the user-
    specified dx and dy keys in model_dict.

    Limitations:
    * The number of distance bins is limited to max_box_distance, defined in the "parameters" function above.

    Arguments:
        * precip
            An iris cube of precipitation to analyse
        * model_dict
            The dictionary for this dataset

    Returns:
        * distance_correlations (max_box_distance):
            Correlations as a function of physical distance, binned by physical distance from the
            central point, using bins delta_x wide.
        * distance_ranges (3, max_box_distance):
            The minimum, median and maximum distance away from the central points, considering
            all points in that distance bin
        * distance_max (max_box_distance):
            The number of bins that contain valid points.  Can be used to subscript
            "distance_correlations" and "distance_ranges" (i.e., distance_correlations[0:distance_max]).
    """


    print('---> Computing correlations for '+str(model_dict['box_size'])+'x'+str(model_dict['box_size'])+' km sub-boxes.')
    longitude = precip.coord('longitude').points
    nlon=len(precip.coord('longitude').points)
    latitude = precip.coord('latitude').points
    nlat=len(latitude)
    max_box_distance,max_boxes,max_timesteps = parameters()
    box_length_x = np.int(model_dict['box_size']//model_dict['dx'])
    box_length_y = np.int(model_dict['box_size']//model_dict['dy'])
#    distance_x = np.zeros(box_length_x)
#    distance_y = np.zeros(box_length_y)
    nboxes=0
    autocorr=0
    npts=np.zeros(max_box_distance,dtype=np.int32)
    distance_lists = np.zeros((max_box_distance,max_box_distance*max_boxes))
    distance_ranges = np.zeros((3,max_box_distance))
    distance_correlations = np.zeros((max_box_distance))

    print('----> Info: Sub-boxes are '+str(box_length_x)+'x'+str(box_length_y)+' gridboxes in this model (nx x ny).')
    for box_xstart in range(0,nlon+1-box_length_x,box_length_x):
        box_xcentre = box_xstart+box_length_x//2
        for box_ystart in range(0,nlat+1-box_length_y,box_length_y):
            box_ycentre = box_ystart+box_length_y//2
            central_precip = precip[:,box_ycentre,box_xcentre]
 #           for box_x in range(box_length_x):
 #               distance_x[box_x]=np.abs(box_xstart+box_x-box_xcentre)*np.cos(latitude[box_ycentre]*np.pi/180.)
 #           for box_y in range(box_length_y):
 #               distance_y[box_y]=np.abs(box_ystart+box_y-box_ycentre)*model_dict['dy']/float(model_dict['dx'])
            for box_x in range(box_length_x):
                for box_y in range(box_length_y):
                    #distance=np.int(round(np.sqrt(distance_x[box_x]*distance_x[box_x]+distance_y[box_y]*distance_y[box_y])))
#                   km_distance=np.sqrt(distance_x[box_x]*distance_x[box_x]+distance_y[box_y]*distance_y[box_y])*model_dict['dx']
                    km_distance = haversine((latitude[box_ycentre],longitude[box_xcentre]),(latitude[box_ystart+box_y],longitude[box_xstart+box_x]))
                    if km_distance < model_dict['box_size']:
                        distance = np.int(round(km_distance/model_dict['dx']))
                        distance_lists[distance,npts[distance]]=km_distance
                        remote_precip=precip[:,box_y+box_ystart,box_x+box_xstart]
                        if (box_x + box_xstart == box_xcentre) and (box_y + box_ystart == box_ycentre):
                            autocorr=autocorr+1
                        corr=np.corrcoef([central_precip.data,remote_precip.data])[1,0]
                        if not np.isnan(corr):
                            distance_correlations[distance]=distance_correlations[distance]+corr
                            npts[distance]=npts[distance]+1
            nboxes = nboxes+1
            if nboxes >= max_boxes :
                raise Exception('ERROR: Number of sub-boxes ('+str(nboxes)+') exceeds maximum number of sub-boxes ('+str(max_boxes)+') exceeded.  Increase size of sub-boxes (-b option) or increase parameter max_boxes in code.')
    distance_max=0
    for dist in range(max_box_distance):
        if npts[dist] > 0 :
            distance_correlations[dist]=distance_correlations[dist]/npts[dist]
            distance_ranges[0,dist]=np.amin(distance_lists[dist,0:npts[dist]])
            distance_ranges[1,dist]=np.median(distance_lists[dist,0:npts[dist]])
            distance_ranges[2,dist]=np.amax(distance_lists[dist,0:npts[dist]])
            distance_max=dist
        else:
            distance_correlations[dist]=-999
        if distance_max==0:
            print("ERROR: no points found within max_box_distance. Try expanding region size.")
    print('---> Info: There are '+str(nboxes)+' sub-boxes in your input data.')
    return(distance_correlations,distance_ranges,distance_max)

def compute_autocorr(precip,model_dict):

    """
    Compute the lagged auto-correlation of precipitatio
    across all points in the analysis domain.

    Arguments:
    * precip:
        An iris cube of precipitation to analyse
    * model_dict:
        The dictionary of information about this dataset

    Returns:
    * time_correlations (max_timesteps):
        Composite lagged auto-correlations across all points
    * time_max (max_timesteps):
        The maximum valid lag in the time_correlation array (can be
        used as a subscript).
    """

    print('---> Computing auto-correlations')
    nlon=len(precip.coord('longitude').points)
    nlat=len(precip.coord('latitude').points)
    autocorr_length = model_dict['autocorr_length']
    max_box_distance,max_boxes,max_timesteps = parameters()
    # +1 to account for lag-zero correlation
    autocorr_nt = np.int64(autocorr_length//(model_dict['dt']))+1
    time_max = autocorr_nt
    time_correlations = np.zeros(max_timesteps)
    print('----> Info: Computing auto-correlations for '+str(autocorr_nt)+' lags.')
    if autocorr_nt > max_timesteps:
        raise Exception('Error: Number of lags for auto-correlation exceeds maximum ('+str(max_timesteps)+').  Increase parameter max_timesteps in code or reduce autocorrelation length.')

    for lon in range(nlon):
        for lat in range(nlat):
            for lag in range(autocorr_nt):
                this_precip = precip.data[:,lat,lon]
                time_correlations[lag] = time_correlations[lag] + np.corrcoef(this_precip,np.roll(this_precip,lag,0))[1,0]

    time_correlations = time_correlations / (nlon*nlat)
    return(time_correlations,time_max)

def compute_spacetime_summary(precip,ndivs,twod=False,cyclic_lon=False,metrics_csv=None,short_name=None,wk_dir='.'):

    """
    Computes summary metrics of spatial and temporal coherence,
    as in Klingaman et al. (2017).

    Method:
    Precipitation data are binned into "ndivs" divisions (at each gridpoint).

    The temporal coherence metric measures the relative frequency of
    persistent upper- and lower-division precipitation to the relative
    frequency of intermittent precipitation (upper-division then
    lower-division, or lower-division then upper-division) on consecutive
    timesteps at the same gridpoint.

    The spatial coherence metric measures the relative frequency of
    persistent upper- and lower-quartile precipitation to the relative