#------------------------------------------------------------------------------
# Name:        01_es_lu_cc.py
# Purpose:     Processing for the CREW project on ES, LUC and CC.
#
# Author:      James Sample
#
# Created:     14/01/2015
# Copyright:   (c) James Sample and JHI, 2015
# License:     https://github.com/JamesSample/ecosystem_services_impacts/blob/master/LICENSE
#------------------------------------------------------------------------------
""" Processes the Future Flows (FF) climate data and estimate climate and land 
    use change effects on Ecosystem Services (ES). Reads workshop outputs and
    performs the following steps:    
    
        1. For each ES, reads monthly rainfall and ET grids for the months
           specified for both baseline and future periods. For the seasons of
           interest, calculates the % change in rainfall and ET between 
           baseline and future.
        
        2. Combines rainfall and runoff percentage changes into a qualitative 
           grid of change in runoff.
           
        3. Estimates impacts grids for each ES for CC only, LUC only and CC &
           LUC combined.

    Inputs grids are supplied in HDF5 file format.
"""

import pandas as pd, h5py, numpy as np, matplotlib, matplotlib.pyplot as plt
import os, sys
from mpl_toolkits.axes_grid1 import ImageGrid
from osgeo import gdal, gdalconst, osr

def read_array_from_h5(h5, variable, model, year, month):
    """ Read an array from a specified location in an H5 file.
    
    Args:
        h5:       The open HDF5 file object
        variable: The variable of interest ('rainfall' or 'pet')
        model:    The code for the climate model of interest (string)
        year:     Year (integer)
        month:    Month (integer)
    
    Returns:
        array
    """
    dset_path = r'/ff_data/%s/%s/%s_%s' % (model, variable, variable, year)
    data = h5.get(dset_path)[:,:,month-1].astype(float)
    
    # Set NoData to NaN
    data[data==-99] = np.nan
    
    # Convert units
    data = data/100

    return data

def avg_rain_et(h5, st_yr, end_yr, months):
    """ Calculate average rainfall and ET grids for the specified years and 
        months.
    
    Args:
        h5:     The open HDF5 file object
        st_yr:  Start year for period of interest (integer)
        end_yr: End year for period of interest (integer)
        months: List of months of interest (integers) 
    
    Returns:
        Tuple of arrays (average rainfall, average PET)
    """
    # Empty arrays to store rainfall and ET totals
    rn_tot = np.zeros((715, 485))
    et_tot = np.zeros((715, 485))
    
    # Total number of years to average over
    years = end_yr + 1 - st_yr
    
    # Loop over rainfall and ET
    for year in range(st_yr, end_yr+1):
        for month in months:
            # Read rainfall and ET grids
            rn = read_array_from_h5(h5, 'rainfall', model, year, month)
            et = read_array_from_h5(h5, 'pet', model, year, month)
            
            # Add to totals
            rn_tot += rn
            et_tot += et
    
    # Average            
    rn_av = rn_tot/years
    et_av = et_tot/years

    return (rn_av, et_av)

def plot_avg_grids(base_rn_av, base_et_av, fut_rn_av, fut_et_av):
    """ Plot the average rainfall and ET grids. Used for testing.
    
    Args:
        base_rn_av: Average rainfall grid for baseline period.
        base_et_av: Average PET grid for baseline period.
        fut_rn_av:  Average rainfall grid for future period.
        fut_et_av:  Average PET grid for future period.
    
    Returns:
        None. Displays maps of each grid using same colour scale.
    """
    # Get min and max values from grids
    rnmin = min(np.nanmin(base_rn_av), np.nanmin(fut_rn_av))    
    rnmax = max(np.nanmax(base_rn_av), np.nanmax(fut_rn_av))
    
    etmin = min(np.nanmin(base_et_av), np.nanmin(fut_et_av))
    etmax = max(np.nanmax(base_et_av), np.nanmax(fut_et_av))
    
    # Plot
    fig = plt.figure()
    grid = ImageGrid(fig, 111,
                     nrows_ncols = (1, 4), 
                     axes_pad=0.5,
                     cbar_mode='each')
    
    im0 = grid[0].imshow(base_rn_av, vmin=rnmin, vmax=rnmax,
                         interpolation='nearest')
    grid.cbar_axes[0].colorbar(im0)
                   
    im1 = grid[1].imshow(fut_rn_av, vmin=rnmin, vmax=rnmax, 
                         interpolation='nearest')
    grid.cbar_axes[1].colorbar(im1)
                                     
    im2 = grid[2].imshow(base_et_av, vmin=etmin, vmax=etmax, 
                         interpolation='nearest')
    grid.cbar_axes[2].colorbar(im2)
                   
    im3 = grid[3].imshow(fut_et_av, vmin=etmin, vmax=etmax, 
                         interpolation='nearest')
    grid.cbar_axes[3].colorbar(im3)
    
    plt.show()

def plot_reclassified_grid(array, out_path, sup_title='Main title', 
                           title='Sub-title'):
    """ Plot and save the reclassified grid.
    
    Args:
        array:     Grid of integers in range -2 to +2
        out_path:  Output file path (PNG or PDF)
        sup_title: Main title for plot (string)
        title:     Sub-title for plot (string)
        
    Returns:
        None. Saves a plot to the specified path.
    """
    # Make a color map of fixed colors
    cmap = matplotlib.colors.ListedColormap(['Red', 'Orange', 'LimeGreen',
                                             'DeepSkyBlue', 'Blue'])

    bounds=[-2.5, -1.5, -0.5, 0.5, 1.5, 2.5]
    norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)
    
    # Create axes for plot (A4 size)
    fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8.3,11.7))
    
    # Plot the array, using the colours specified
    img = axes.imshow(array, interpolation='nearest', origin='upper',
                      cmap=cmap, norm=norm)

    # Add labels to plot
    plt.title(title)
    plt.suptitle(sup_title, fontsize=16, y=0.95)
    plt.ylabel('Northing')
    plt.xlabel('Easting')
    plt.grid(True)

    # Reformat the axis labels (mainly change the Y values into northings)
    axes.set_yticks([35, 135, 235, 335, 435, 535, 635, 735])
    axes.set_yticklabels([1200, 1100, 1000, 900, 800, 700, 600, 500])
    axes.set_xticks([100, 200, 300, 400])

    # Add axes for the color bar
    cax = fig.add_axes([0.2, 0.785, 0.02, 0.10])

    # Add the colour bar and set labels
    cbar = fig.colorbar(img, cax=cax, cmap=cmap, norm=norm, boundaries=bounds,
                        ticks=[-2.2,-1.2,-0.2,0.8,1.8])

    cbar.set_ticklabels(['Large decrease',
                         'Small decrease',
                         'Neutral',
                         'Small increase',
                         'Large increase'], update_ticks=True)

    # Make the cbar ticks invisible
    ticks = cbar.ax.get_yticklines()
    for tick in ticks:
        plt.setp(tick, alpha=0)

    cbar_labels = plt.getp(cbar.ax.axes, 'yticklabels')
    plt.setp(cbar_labels, fontsize=10) 

    # Save fig
    plt.savefig(out_path, dpi=300)
##    plt.show()
    plt.clf()
    plt.close()

def reclass_rn_et_grid(array):
    """ Take an array of percentage changes and reclassify it according to:
    
            % change  | Class
             x<=-15   |  -2
            -15<x<=-5 |  -1
            -5<x<=5   |   0
             5<x<=15  |  +1
             15<x     |  +2 
    
    Args:
        array: Array of percentage changes to be reclassified.
        
    Returns:
        Reclassified array
    """
    # Create copy of array for reclass values         
    rc = array.copy()
    rc[array<=-15] = -2
    rc[(-15<array) & (array<=-5)] = -1
    rc[(-5<array) & (array<=5)] = 0
    rc[(5<array) & (array<=15)] = 1
    rc[15<array] = 2
    
    return rc

def reclass_ro(matrix_path, rn, et):
    """ Generate reclassification matrix for runoff based on reclassified 
        change grids for rainfall and PET and the runoff reclassification 
        matrix from the workshop.
    
    Args:
        matrix_path: Path to CSV file representing runoff matrix.
        rn:          Reclassified rainfall grid from reclass_rn_et_grid
        et:          Reclassified PET grid from reclass_rn_et_grid
    
    Returns:
        Array (grid of integers representing change in runoff)
    """
    # Read matrix
    df = pd.read_csv(matrix_path, index_col=0)
    
    # Grid of NaNs wih correct shape
    ro = rn.copy()*np.nan
    
    # Loop over inidces
    for x, y in np.ndindex(ro.shape):
        # Get values for change in rainfall and ET
        et_ch = et[x, y]
        rn_ch = rn[x, y]
        
        # If both are not nan, reclassify
        if (np.isfinite(et_ch) and np.isfinite(rn_ch)):
            rc_val = df.ix[int(et_ch), str(int(rn_ch))]
            ro[x, y] = rc_val
    
    return ro

def reclass_es_ro(es_idx, ro):
    """ Reclassify the runoff grid to estimate effects of runoff change on each
        ES.
    
    Args:
        es_idx: The ID of the ES of interest in data frame ro_df
        ro:     The runoff change grid from reclass_ro
        
    Returns:
        Array (grid of integers representing change in ES) 
    """
    # Make a copy of the ro grid to update
    es = ro.copy()
    
    # Reclassify
    for chng in [-2, -1, 0, 1, 2]:
        es[ro==chng] = ro_df.ix[es_idx, 'RO_%d' % chng]
    
    return es

def read_ascii(ascii_path,
               xmin=0,
               xmax=485000,
               ymin=520000,
               ymax=1235000,
               exptd_rows=715,
               exptd_cols=485,
               exptd_px_wd=1000,
               exptd_px_ht=-1000,
               exptd_ndv=-9999):
    """ Read an ASCII grid file, clip it to the specified bounding box and
        return a numpy array.
    
    Args:
        xmin:         Minimum Easting in OSGB1936 metres.
        xmax:         Maximum Easting in OSGB1936 metres.
        ymin:         Minimum Northing in OSGB1936 metres.
        ymax:         Maximum Northing in OSGB1936 metres.
        exptd_rows:   No. of rows expected in file.
        exptd_cols:   No. of columns expected in file.
        exptd_px_wd:  Cell width.
        exptd_px_ht:  Cell height.
        exptd_ndv:    No data value.
        
    Returns:
        Array (floats).
    """
    # Register drivers
    gdal.AllRegister()

    # Process the file with GDAL
    ds = gdal.Open(ascii_path, gdalconst.GA_ReadOnly)
    if ds is None:
    	print 'Could not open ' + ascii_path
    	sys.exit(1)

    # In order to select the first cell correctly, choose a point just within
    # the top left corner of the specified bounding box.
    x = xmin + 10
    y = ymax - 10

    # Dataset properties
    geotransform = ds.GetGeoTransform()
    originX = geotransform[0]
    originY = geotransform[3]
    pixelWidth = geotransform[1]
    pixelHeight = geotransform[5]

    # Calculate number of rows and cols to return
    rows = abs(int((ymax-ymin)/pixelHeight))
    cols = int((xmax-xmin)/pixelWidth)

    # Select starting pixel
    xOffset = int((x - originX) / pixelWidth)
    yOffset = int((y - originY) / pixelHeight)

    band = ds.GetRasterBand(1)
    no_data_val = band.GetNoDataValue()

    # Simple checking
    assert rows == exptd_rows
    assert cols == exptd_cols
    assert pixelWidth == exptd_px_wd
    assert pixelHeight == exptd_px_ht
    assert no_data_val == exptd_ndv

    # Read the data to an array
    data = band.ReadAsArray(xOffset, yOffset, cols, rows)

    # Close the dataset
    ds = None

    return data.astype(float)

def process_land_use_change(lu_mat_path, base, fut, esid, codes_df):
    """ Estimate land use change (LUC) only effects for the specified ES.
    
    Args:
        lu_mat_path: Excel file containing land use matrices from the workshop.
        base:        Baseline land luse grid.
        fut:         Future land luse grid.
        esid:        ES ID from land use matrices Excel file
        codes_df:    Land use code look-up table (as data frame)
        
    Returns:
        Array (grid of integers representing change in ES)  
    """
    # Read matrix for this ES
    lu_mat = pd.read_excel(lu_mat_path, sheetname='Land Use')

    # Get row for start of matrix
    st_row = (lu_mat['ES_ID']==esid).nonzero()[0][0] + 2

    # Read matrix of interest
    lu_mat = pd.read_excel(lu_mat_path, sheetname='Land Use', skiprows=st_row,
                           skip_footer=(120-6-st_row), parse_cols='C:I',
                           index_col=0)

    # Perform reclassification
    # Grid of NaNs wih correct shape
    rc = base.copy()*np.nan
    
    # Loop over inidces
    for x, y in np.ndindex(base.shape):
        # Get values for baseline and future LU
        base_lu = base[x, y]
        fut_lu = fut[x, y]
        
        # If both are not nan, reclassify
        if (np.isfinite(base_lu) and np.isfinite(fut_lu)):
            # Get the base and fut LU as a string
            base_str = codes_df.ix[int(base_lu)]['LU_Class']
            fut_str = codes_df.ix[int(fut_lu)]['LU_Class']
            rc_val = lu_mat.ix[base_str, fut_str]
            rc[x, y] = rc_val
    
    return rc

def process_land_use_and_climate_change(lucc_mat_path, lugrid, ccgrid, esid):
    """ Estimate combined land use and climate change effects for the specified
        ES.
        
    Args:
        lucc_mat_path: Excel file containing matrices from the workshop.
        lugrid:        The grid of land use change effects.
        ccgrid:        The grid of climate change effects.
        esid:          ES ID from workshop matrices Excel file.
        
    Returns:
        Array (grid of integers representing change in ES) 
    """
    # Read matrix for this ES
    lucc_mat = pd.read_excel(lucc_mat_path, sheetname='CC_LU')

    # Get row for start of matrix
    st_row = (lucc_mat['ES_ID']==esid).nonzero()[0][0] + 2
    
    # Read matrix of interest
    lucc_mat = pd.read_excel(lucc_mat_path, sheetname='CC_LU', skiprows=st_row,
                             skip_footer=(108-5-st_row), parse_cols='C:I',
                             index_col=0)

    # Perform reclassification
    # Grid of NaNs wih correct shape
    rc = lugrid.copy()*np.nan
    
    # Loop over inidces
    for x, y in np.ndindex(lugrid.shape):
        # Get values for baseline and future LU
        lu = lugrid[x, y]
        cc = ccgrid[x, y]
        
        # If both are not nan, reclassify
        if (np.isfinite(lu) and np.isfinite(cc)):
            # Get the base and fut LU as a string
            rc_val = lucc_mat.ix[int(lu), int(cc)]
            rc[x, y] = rc_val
    
    return rc

def array_to_gtiff(out_path, data_array, ndv=-9999, xmin=0, ymax=1235000, 
                   cell_size=1000):
    """ Convert numpy array to 16-bit integer GeoTiff.
    
    Args:
        out_path:   The .tif file to be created.
        data_array: The (integer) data array to save.
        ndv:        No data value.
        xmin:       Minimum x (Easting) co-ordinate, in OSGB1936 metres
        ymax:       Maximim y (Northing) co-ordinate, in OSGB1936 metres
        cell_size:  Cell size (metres)
    
    Returns:
        None. Array is saved to specified path.
    """
    # Copy data_array so that it is not modified
    data = data_array.copy()
    
    # Convert NaNs to NDV
    data[np.isnan(data)] = ndv
    
    # Get array shape
    cols = data.shape[1]
    rows = data.shape[0]

    # Get driver
    driver = gdal.GetDriverByName('GTiff') # NB can't directly create ArcInfo ASCII grids in this way

    # Create a new raster data source
    out_ds = driver.Create(out_path, cols, rows, 1, gdal.GDT_Int16)

    # Get spatial ref details
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(27700) # From EPSG for OSGB36 grid

    # Write metadata
    out_ds.SetGeoTransform((xmin, cell_size, 0.0, ymax, 0.0, -1*cell_size)) #(xmin, cellsize, 0, ymax, 0, -cellsize)
    out_ds.SetProjection(srs.ExportToWkt())
    out_band = out_ds.GetRasterBand(1)
    out_band.SetNoDataValue(ndv)
    out_band.WriteArray(data)

    # Tidy up
    del out_ds, out_band
    
# #############################################################################
# User input
# Climate data
ff_h5_path = r'D:\WBM_Development_2014\WBM_2014_Monthly_Input_File.h5'

# Runoff matrices
ro_path = r'D:\Eco_Services_Impacts\Matrices_Development\03_Group_1_Matrices\Runoff_Impacts_Grp1.csv'
ro_matrix_15 = r'D:\Eco_Services_Impacts\Matrices_Development\02_Common_Matrices\Runoff_Matrix_15pct.csv'

# Land use data
base_path = r'D:\Eco_Services_Impacts\Land_Use\baseline_lu_lcm07.txt'
fut_path = r'D:\Eco_Services_Impacts\Land_Use\future_lu_2050.txt'

# Land use matrices
lu_classes_path = r'D:\Eco_Services_Impacts\Land_Use\Land_Use_Classes.csv'
lu_matrices_path = r'D:\Eco_Services_Impacts\Matrices_Development\03_Group_1_Matrices\Land_Use_Matrices_Grp1.xlsx'

# Land use and climate combined matrices
lucc_matrices_path = r'D:\Eco_Services_Impacts\Matrices_Development\03_Group_1_Matrices\Climate_And_Land_Use_Matrices_Grp1.xlsx'

# Output folders
out_pdf_fold = r'D:\Eco_Services_Impacts\Model_Output\02_Group_1_Output\PDF'
out_array_fold = r'D:\Eco_Services_Impacts\Model_Output\02_Group_1_Output\GeoTiffs'

# Time periods to compare
base_st_yr, base_end_yr = 1961, 1990
fut_st_yr, fut_end_yr = 2041, 2070

# Future Flows models of interest
models = ['afixa', 'afixc', 'afixl', 'afixm', 'afixo', 'afixh', 
          'afixi', 'afixj', 'afixk', 'afgcx', 'afixq']
# #############################################################################

# Read LU grids
base = read_ascii(base_path)
base[base==-9999] = np.nan
fut = read_ascii(fut_path)
fut[fut==-9999] = np.nan

# Read LU class codes
codes_df = pd.read_csv(lu_classes_path, index_col=0)

# Read the runoff matrices
ro_df = pd.read_csv(ro_path, index_col=0)

# Open H5 file
h5 = h5py.File(ff_h5_path, 'r')

# Iterate over each ES
for idx in ro_df.index:
    print '\nProcessing land use change impacts for %s.' % ro_df.ix[idx, 'ES']
    
    # 1. Process land use change only
    luc = process_land_use_change(lu_matrices_path, base, fut, idx, codes_df)
    
    # Prepare to save
    out_name = 'ES%02d_LUC' % idx
    
    # Save array
    out_array = os.path.join(out_array_fold, '%s.tif' % out_name)
    array_to_gtiff(out_array, luc)
    
    # Save PDF
    out_pdf = os.path.join(out_pdf_fold, '%s.pdf' % out_name)
    plot_reclassified_grid(luc, out_pdf,
                           sup_title='Change in %s' % ro_df.ix[idx, 'ES'],
                           title='(land use change only)' )

    # 2. Process climate change only    
    # Get the relevant months for this ES
    months = [int(i) for i in ro_df.ix[idx, 'Key_Months'].split(',')]

    # Loop over climate models of interest
    for model in models:
        print ('Processing climate change impacts for '
               '%s (model %s).' % (ro_df.ix[idx, 'ES'], model))
        
        # 2.1. Baseline
        base_rn_av, base_et_av = avg_rain_et(h5, base_st_yr, base_end_yr,
                                             months)

        # 2.2. Future
        fut_rn_av, fut_et_av = avg_rain_et(h5, fut_st_yr, fut_end_yr,
                                           months)

        # Plot
#        plot_avg_grids(base_rn_av, base_et_av, fut_rn_av, fut_et_av)
        
        # Calculate % change
        rn_pct = 100*(fut_rn_av - base_rn_av)/base_rn_av
        et_pct = 100*(fut_et_av - base_et_av)/base_et_av

        # Reclassify
        rn_rc = reclass_rn_et_grid(rn_pct)
        et_rc = reclass_rn_et_grid(et_pct)

#        plot_reclassified_grid(rn_rc)
#        plot_reclassified_grid(et_rc)
        
        # Generate runoff grid
        ro = reclass_ro(ro_matrix_15, rn_rc, et_rc)
        
#        # Plot runoff grid        
#        plot_reclassified_grid(ro, 
#                               sup_title='Change in runoff',
#                               title='(Model %s; %s)' % (model, months))
        
        # Reclass ro grid to estimate ES impact
        es = reclass_es_ro(idx, ro)

        # Prepare to save
        out_name = 'ES%02d_%s' % (idx, model)
        
        # Save array
        out_array = os.path.join(out_array_fold, '%s.tif' % out_name)
        array_to_gtiff(out_array, es)
        
        # Save PDF
        out_pdf = os.path.join(out_pdf_fold, '%s.pdf' % out_name)   
        plot_reclassified_grid(es, out_pdf,
                               sup_title='Change in %s' % ro_df.ix[idx, 'ES'],
                               title='(climate model %s only)' % model)
        
        # 3. Process combined land use and climate effects
        print ('Processing climate and land use change impacts for '
               '%s (model %s).' % (ro_df.ix[idx, 'ES'], model))
               
        # Reclassify to get CC and LUC effects
        cc_lu = process_land_use_and_climate_change(lucc_matrices_path, luc,
                                                    es, idx)

        # Prepare to save
        out_name = 'ES%02d_LUC_%s' % (idx, model)
        
        # Save array
        out_array = os.path.join(out_array_fold, '%s.tif' % out_name)
        array_to_gtiff(out_array, cc_lu)
        
        # Save PDF
        out_pdf = os.path.join(out_pdf_fold, '%s.pdf' % out_name)   
        plot_reclassified_grid(cc_lu, out_pdf,
                               sup_title='Change in %s' % ro_df.ix[idx, 'ES'],
                               title='(climate and land use change together)')
                               
# Close H5 file
h5.close()

print '\nFinished.'            