'''
Created on 21/02/2017
@author: ramato

Script plots precipitation data from bigendian, rim removed, 30 year averaged monthly files.
Converts to mm/day, creates a plot for each month (6 year average) on log / linear scales, or both.

In rid_dic (line ~25), 'raw data' / 'difference' dictates how the plot looks.
To run this code fewer datasets, comment out lines in rid_dic.
To run on different data (first prep to match line 5 description):
    amend 'DIR', 'YEARS' 'RIS_DIC' and 'fil' (in load_file()).
'''
import iris
import calendar
import sys
import numpy as np
import math
import iris.plot as iplt
import matplotlib.pyplot as plt
import matplotlib.colors as mcols
import cartopy.crs as ccrs # coordinate reference system
import iris.coord_categorisation as coord_cat

iris.FUTURE.netcdf_promote = True
#pylint: disable=no-member
#pylint: disable=trailing-whitespace
MONTHS = calendar.month_abbr

#DIR = '/home/h03/ramato/CSSP_validation'
DIR = '/project/precis/rosie/CSSP'
YEARS = '1980-1985'
RID_DIC = {#'apepba' : ['EraInt', 'raw data'],
           'apepca' : ['20CR', 'raw data'],
           #'apepb-apepc' : ['ERAInt-20CR', 'difference'],
           }

COL_DIC = {'precipitation_flux' : ['Spectral', 'YlGnBu'],
           'air_temperature' : ['Spectral_r', 'Spectral_r'],
           }

def rounddown(num):
    '''Round a number to it's lower divisor by 10 (e.g 57 to 50),
    unless that rounds to 0; in which case round to an integar.'''
    if num > 10:
        return int(math.floor(num / 10.0)) * 10
    else:
        return math.floor(num)

def load_file(rid):
    '''Load  in monthly files to a single cube'''
    #fil = '{}/ERA_20CR_validation/bigend/{}.h9i5*.05216.be_norim.pp'.format(DIR, rid)
    #fil = '{}/ERA_20CR_validation/bigend/{}.h9i5*.03236.128.be_norim.pp'.format(DIR, rid)
    fil = '{}/monitoring/data/20CR.1975.2009.nc'.format(DIR)
    cube = iris.load_cube(fil, 'precipitation_flux')
    return cube

def mmday(cube):
    '''Convert precip data from kg/m/s to mm/day'''
    new_cube = cube * 86400
    new_cube.units = 'mm/day'
    new_cube.standard_name = 'precipitation_flux'
    return new_cube

def get_mon(i):
    if i>0:
        mon = MONTHS[i]
    elif i == 0:
        mon = MONTHS[12]#dec
    return mon

def plot_subplot(cube, i, mycmap, myrange, mon):
    ax1 = plt.subplot(4, 3, i, projection=ccrs.PlateCarree())
    ax1.coastlines()
    #ax1.set_extent([-45, 55, 20, 80], ccrs.PlateCarree())
    mesh = iplt.contourf(cube, cmap=mycmap, levels=myrange)
    plt.title('{}'.format(mon))
    add_gridlines(ax1)
    return mesh

def make_colourbar(mesh, myextend, tick_levels, cbar_title):
    '''standard setup for plot colourbar'''
    cax = plt.axes([0.1, 0.1, 0.8, 0.025]) 
    cbar = plt.colorbar(mesh, cax=cax, orientation='horizontal', extend=myextend)
    cbar.set_ticks(tick_levels)
    cbar.set_label(cbar_title)

def add_gridlines(axs):
    '''Add coastlines and gridlines, showing lat-lon values'''
    axs.coastlines()
    gls = axs.gridlines(draw_labels=True)
    gls.xlabels_top = None
    gls.ylabels_right = None

def plot_monthly_precip(cube, model, plot_type, rid):
    '''Create spatial plots of precipitation data with a logarithmic colourscale;
    this is best for highlighting areas of extreme precip (changes). 
    Model output data (raw data) is plotted on a scale from 0 to max.
    Difference data (model1 - model2) is plotted on a scale centred at zero.'''
    # Define scaling levels for the logarithmic colouring.
    maximum_scale_level = max(np.max(cube.data), abs(np.min(cube.data)))
    mymax = rounddown(maximum_scale_level)
    #mymax = 125

    if plot_type == 'difference':
        myextend = 'both'
        mymin = -mymax
        tick_levels = np.linspace(mymin, mymax, 9)
        mycmap = plt.get_cmap(COL_DIC.get(cube.standard_name)[0])
        #tick_levels = [-mymax, -mymax/5, -mymax/10, 0.0, mymax/10, mymax/5, mymax]      
    elif plot_type == 'raw data':
        myextend = 'max'
        if cube.units == 'mm/day':
            mymin = rounddown(np.min(cube.data))
        elif cube.units == 'celsius':
            mymin = -mymax
        print mymin, np.min(cube.data)
        tick_levels = np.linspace(mymin, mymax, 5)
        #tick_levels = [0.0, mymax/10, mymax/5, mymax]
        key = cube.standard_name        
        col_scale =  COL_DIC.get(cube.standard_name)[1]
        print key, col_scale
        mycmap = plt.get_cmap(COL_DIC.get(cube.standard_name)[1])

    fig = plt.figure(figsize=(16, 10))
    fig.text(0.5, 0.97, '{} {} \n ({} average)'.format(model, cube.standard_name, YEARS),
             horizontalalignment='center', verticalalignment='top',
             fontsize=15, fontweight='bold') 
    cbar_title = '{} ({})'.format(cube.standard_name, cube.units)
    myrange = np.arange(mymin, mymax+1)

    for i in range(0, cube.shape[0]):
        mon = cube.coord('month')[i].points[0]
        print i, mon
        mesh = plot_subplot(cube[i], i, mycmap, myrange, mon)
    make_colourbar(mesh, myextend, tick_levels, cbar_title)
    #plt.savefig('{}/linear_plots/{}.monthly.h9i5.03236.128.png'.format(DIR, rid))
    plt.show()


def main():
    '''Run the script from here. To remove e.g. linear scale plots, comment out that function here. 
    Test with one month by amending months[1:] to e.g. months[1:2] for January (line ~157)'''
    for _, key in enumerate(RID_DIC):
        rid = key
        model = RID_DIC.get(key)[0]
        plot_type = RID_DIC.get(key)[1]
        print rid, model

        cube = load_file(rid)
        iris.coord_categorisation.add_month(cube, 'time', name='month')
        cube = cube.aggregated_by(['month'], iris.analysis.MEAN)        
        print cube
        #print cube[1]

        #if cube.standard_name == 'precipitation_flux':
        #    cube = mmday(cube)
        #elif cube.units == 'K':
        #    cube.convert_units('celsius')
        #else:
        #    print 'which cube variable??'
        plot_monthly_precip(cube, model, plot_type, rid)
        print cube

if __name__ == "__main__":
    main()