#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jul 31 15:38:34 2019

@author: zoescrewvala
"""

import os
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
import pandas as pd
import xarray as xr
import pygemfxns_gcmbiasadj as pygemfxns

#%%
def avg_parameters_20yr(component, area):
    """
    sum up annually and convert to km^3, average every 20 yrs, then convert to m w.e. and average every 20 yrs again
    """
    component_annual = ((pygemfxns.annual_sum_2darray(component))/1000)*area
    component_annual_regional_1 = np.sum(component_annual[1:], axis=0) # in km^3 aka gigatons
    component_20yr_gt = np.add.reduceat(component_annual_regional_1, np.arange(0, len(component_annual_regional_1), 20))
#    component_avg_gt[i] = component_20yr
    component_annual_regional_2 = ((component_annual_regional_1)/(np.sum(area, axis=0)))*1000 # km^3 (gigatons) --> m w.e.
    component_20yr_mwea = np.add.reduceat(component_annual_regional_2, np.arange(0, len(component_annual_regional_2), 20))
#    component_avg_mwea[i] = component_20yr

    return component_annual_regional_1, component_annual_regional_2, component_20yr_gt, component_20yr_mwea
#
#def create_avg_arrays(component):
#    """
#    returns arrays of zeros in correct shape to fill with averages later
#    """
#    component_tide = np.zeros((18,1452), dtype=float)
#    component_nontide = np.zeros((98,1452), dtype=float)
#    
#    component_avg_gt = np.zeros((10,6), dtype=float)
#    component_avg_mwea = np.zeros((10,6), dtype=float)
#    component_tide_avg_gt = np.zeros((10,6), dtype=float)
#    component_tide_avg_mwea = np.zeros((10,6), dtype=float)
#    component_nontide_avg_gt = np.zeros((10,6), dtype=float)
#    component_nontide_avg_mwea = np.zeros((10,6), dtype=float)
#    
#    return component_tide, component_nontide, component_avg_gt, component_avg_mwea, component_tide_avg_gt, component_tide_avg_mwea, component_nontide_avg_gt, component_nontide_avg_mwea

def avg_all_params(component1, component2, component3, component4):
    avg_1d_1 = np.average(component1, axis=0)
    avg_1d_2 = np.average(component2, axis=0)
    avg_1d_3 = np.average(component3, axis=0)
    avg_1d_4 = np.average(component4, axis=0)
    
    return avg_1d_1, avg_1d_2, avg_1d_3, avg_1d_4


#%%
label_list = ['Average (all glaciers)', 'Tidewater', 'Non-tidewater']
sim_list = ['CanESM2', 'CCSM4', 'CSIRO-Mk3-6-0', 'CNRM-CM5', 'GFDL-CM3', 'GFDL-ESM2M', 'GISS-E2-R', 'IPSL-CM5A-LR', 
            'MPI-ESM-LR', 'NorESM1-M']
rcp_list = ['26', '45', '85']
RCP_list = ['RCP 2.6', 'RCP 4.5', 'RCP 8.5']
color_list = ['b', 'm', 'r', 'k']

# tidewater glacier indexes from larsen
tidewater_idxs = [13, 26, 36, 46, 52, 54, 55, 70, 71, 72, 73, 87, 90, 92, 93, 96, 109, 113]

# prep for averaging gcms
acc_tide, ref_tide, mlt_tide, abl_tide = np.zeros((18,1452), dtype=float), np.zeros((18,1452), dtype=float), np.zeros((18,1452), dtype=float), np.zeros((18,1452), dtype=float)
acc_nontide, ref_nontide, mlt_nontide, abl_nontide = np.zeros((98,1452), dtype=float), np.zeros((98,1452), dtype=float), np.zeros((98,1452), dtype=float), np.zeros((98,1452), dtype=float)
acc_avg_gt, ref_avg_gt, mlt_avg_gt, abl_avg_gt = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)
acc_avg_mwea, ref_avg_mwea, mlt_avg_mwea, abl_avg_mwea = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)
acc_tide_avg_gt, ref_tide_avg_gt, mlt_tide_avg_gt, abl_tide_avg_gt = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)
acc_tide_avg_mwea, ref_tide_avg_mwea, mlt_tide_avg_mwea, abl_tide_avg_mwea = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)
acc_nontide_avg_gt, ref_nontide_avg_gt, mlt_nontide_avg_gt, abl_nontide_avg_gt = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)
acc_nontide_avg_mwea, ref_nontide_avg_mwea, mlt_nontide_avg_mwea, abl_nontide_avg_mwea = np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float), np.zeros((10,6), dtype=float)

area_tide = np.zeros((18,121), dtype=float)
area_nontide = np.zeros((98,121), dtype=float)

for j in range(len(rcp_list)):  
    fig, axes = plt.subplots(2, 3, squeeze=False, sharex='col', sharey='row', 
                 gridspec_kw = {'wspace':0.4, 'hspace':0.15})
    for i in range(len(sim_list)):
        ds = xr.open_dataset(os.getcwd() + '/../Output/simulations/' + sim_list[i] + '/R1_' + sim_list[i] + '_rcp' + rcp_list[j] + '_c1_ba1_1sets_1980_2100.nc')
        time = ds.variables['year'].values[:]
        area = ds.variables['area_glac_annual'].values[:,:-1,0]
#        massbal = ds.variables['massbaltotal_glac_monthly'].values[:,:,0]
        acc = ds.variables['acc_glac_monthly'].values[:,:,0]
        ref = ds.variables['refreeze_glac_monthly'].values[:,:,0]
        mlt = ds.variables['melt_glac_monthly'].values[:,:,0]
        abl = ds.variables['frontalablation_glac_monthly'].values[:,:,0]

        # separate tidewater and non-tidewater glaciers
        nontidewater_idxs = 
        
        acc_tide = acc[tidewater_idxs, :]
        ref_tide = ref[k,:]
        mlt_tide = mlt[k,:]
        abl_tide = abl[k,:]
        area_tide = area[k,:]
        
        row_t = 0
        row_nt = 0
        for k in range(116):
            if k in tidewater_idxs:
                acc_tide[row_t] = acc[k,:]
                ref_tide[row_t] = ref[k,:]
                mlt_tide[row_t] = mlt[k,:]
                abl_tide[row_t] = abl[k,:]
                area_tide[row_t] = area[k,:]

                row_t += 1
            else:
                acc_nontide[row_t] = acc[k,:]
                ref_nontide[row_t] = ref[k,:]
                mlt_nontide[row_t] = mlt[k,:]
                abl_nontide[row_t] = abl[k,:]
                area_nontide[row_t] = area[k,:]

                row_nt += 1

        # 20 yr avg for ALL glaciers
        acc_annual_regional_1, acc_annual_regional_2, acc_avg_gt[i], acc_avg_mwea[i] = avg_parameters_20yr(acc, area)
        ref_annual_regional_1, ref_annual_regional_2, ref_avg_gt[i], ref_avg_mwea[i] = avg_parameters_20yr(ref, area)
        mlt_annual_regional_1, mlt_annual_regional_2, mlt_avg_gt[i], mlt_avg_mwea[i] = avg_parameters_20yr(mlt, area)
        abl_annual_regional_1, abl_annual_regional_2, abl_avg_gt[i], abl_avg_mwea[i] = avg_parameters_20yr(abl, area)

        acc_annual_regional_1t, acc_annual_regional_2t, acc_tide_avg_gt[i], acc_tide_avg_mwea[i] = avg_parameters_20yr(acc_tide, area_tide)
        ref_annual_regional_1t, ref_annual_regional_2t, ref_tide_avg_gt[i], ref_tide_avg_mwea[i] = avg_parameters_20yr(ref_tide, area_tide)
        mlt_annual_regional_1t, mlt_annual_regional_2t, mlt_tide_avg_gt[i], mlt_tide_avg_mwea[i] = avg_parameters_20yr(mlt_tide, area_tide)
        abl_annual_regional_1t, abl_annual_regional_2t, abl_tide_avg_gt[i], abl_tide_avg_mwea[i] = avg_parameters_20yr(abl_tide, area_tide)

        acc_annual_regional_1nt, acc_annual_regional_2nt, acc_nontide_avg_gt[i], acc_nontide_avg_mwea[i] = avg_parameters_20yr(acc_nontide, area_nontide)
        ref_annual_regional_1nt, ref_annual_regional_2nt, ref_nontide_avg_gt[i], ref_nontide_avg_mwea[i] = avg_parameters_20yr(ref_nontide, area_nontide)
        mlt_annual_regional_1nt, mlt_annual_regional_2nt, mlt_nontide_avg_gt[i], mlt_nontide_avg_mwea[i] = avg_parameters_20yr(mlt_nontide, area_nontide)
        abl_annual_regional_1nt, abl_annual_regional_2nt, abl_nontide_avg_gt[i], abl_nontide_avg_mwea[i] = avg_parameters_20yr(abl_nontide, area_nontide)
        
        ds.close()
    
    
                                    ### plotting average all glaciers ###
    
    # averaging parameters
    acc_avg_gt_1d, ref_avg_gt_1d, mlt_avg_gt_1d, abl_avg_gt_1d = avg_all_params(acc_avg_gt, ref_avg_gt, mlt_avg_gt, abl_avg_gt)
    acc_avg_mwea_1d, ref_avg_mwea_1d, mlt_avg_mwea_1d, abl_avg_mwea_1d = avg_all_params(acc_avg_mwea, ref_avg_mwea, mlt_avg_mwea, abl_avg_mwea)
#    np.average()
    
    acc_tide_avg_gt_1d, ref_tide_avg_gt_1d, mlt_tide_avg_gt_1d, abl_tide_avg_gt_1d = avg_all_params(acc_tide_avg_gt, ref_tide_avg_gt, mlt_tide_avg_gt, abl_tide_avg_gt)
    acc_tide_avg_mwea_1d, ref_tide_avg_mwea_1d, mlt_tide_avg_mwea_1d, abl_tide_avg_mwea_1d = avg_all_params(acc_tide_avg_mwea, ref_tide_avg_mwea, mlt_tide_avg_mwea, abl_tide_avg_mwea)
    
    acc_nontide_avg_gt_1d, ref_nontide_avg_gt_1d, mlt_nontide_avg_gt_1d, abl_nontide_avg_gt_1d = avg_all_params(acc_nontide_avg_gt, ref_nontide_avg_gt, mlt_nontide_avg_gt, abl_nontide_avg_gt)
    acc_nontide_avg_mwea_1d, ref_nontide_avg_mwea_1d, mlt_nontide_avg_mwea_1d, abl_nontide_avg_mwea_1d = avg_all_params(acc_nontide_avg_mwea, ref_nontide_avg_mwea, mlt_nontide_avg_mwea, abl_nontide_avg_mwea)
    
    x_values = time[0:120:20]
    y_values = acc_avg_gt_1d
    y2_values = ref_avg_gt_1d
    y3_values = np.negative(mlt_avg_gt_1d)
    y4_values = np.negative(abl_avg_gt_1d)
    net = y_values + y2_values + y3_values + y4_values
    
    axes[0,0].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[0,0].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[0,0].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[0,0].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[0,0].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[0,0].set_title(label='Average (all glaciers)', size=10, horizontalalignment='center', verticalalignment='baseline')

    y_values = acc_tide_avg_gt_1d
    y2_values = ref_tide_avg_gt_1d
    y3_values = np.negative(mlt_tide_avg_gt_1d)
    y4_values = np.negative(abl_tide_avg_gt_1d)
    net = y_values + y2_values + y3_values + y4_values
    
    axes[0,1].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[0,1].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[0,1].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[0,1].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[0,1].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[0,1].set_title(label='Tidewater', size=10, horizontalalignment='center', verticalalignment='baseline')
    
    y_values = acc_nontide_avg_gt_1d
    y2_values = ref_nontide_avg_gt_1d
    y3_values = np.negative(mlt_nontide_avg_gt_1d)
    y4_values = np.negative(abl_nontide_avg_gt_1d)
    net = y_values + y2_values + y3_values + y4_values

    axes[0,2].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[0,2].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[0,2].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[0,2].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[0,2].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[0,2].set_title(label='Non-tidewater', size=10, horizontalalignment='center', verticalalignment='baseline')
    
    
    ### in m w.e. ###
    y_values = acc_avg_mwea_1d
    y2_values = ref_avg_mwea_1d
    y3_values = np.negative(mlt_avg_mwea_1d)
    y4_values = np.negative(abl_avg_mwea_1d)
    net = y_values + y2_values + y3_values + y4_values
    
    axes[1,0].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[1,0].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[1,0].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[1,0].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[1,0].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[1,0].set_title(label='Average (all glaciers)', size=10, horizontalalignment='center', verticalalignment='baseline')

    y_values = acc_tide_avg_mwea_1d
    y2_values = ref_tide_avg_mwea_1d
    y3_values = np.negative(mlt_tide_avg_mwea_1d)
    y4_values = np.negative(abl_tide_avg_mwea_1d)
    net = y_values + y2_values + y3_values + y4_values
    
    axes[1,1].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[1,1].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[1,1].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[1,1].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[1,1].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[1,1].set_title(label='Tidewater', size=10, horizontalalignment='center', verticalalignment='baseline')
    
    y_values = acc_nontide_avg_mwea_1d
    y2_values = ref_nontide_avg_mwea_1d
    y3_values = np.negative(mlt_nontide_avg_mwea_1d)
    y4_values = np.negative(abl_nontide_avg_mwea_1d)
    net = y_values + y2_values + y3_values + y4_values

    axes[1,2].bar(x_values, y_values, color='b', width=12, zorder=2)
    axes[1,2].bar(x_values, y2_values, bottom=y_values, color='m', width=12, zorder=3)
    axes[1,2].bar(x_values, y3_values, color='r', width=12, zorder=2)
    axes[1,2].bar(x_values, y4_values, bottom=y3_values, color='k', width=12, zorder=3)
    axes[1,2].scatter(x_values, net, color='c', marker='_', s=16, zorder=5)
    axes[1,2].set_title(label='Non-tidewater', size=10, horizontalalignment='center', verticalalignment='baseline')

    #figure style
    for i in range(3):
        axes[1,i].set_xlabel('Year', size=12)
        axes[1,i].xaxis.set_major_locator(plt.MultipleLocator(40))
        axes[1,i].xaxis.set_minor_locator(plt.MultipleLocator(10))
        axes[1,i].tick_params(axis='x', rotation=70)
    
    # Y-label
    axes[0,0].set_ylabel('Mass balance [Gt yr^-1]', size=10)
    axes[1,0].set_ylabel('Mass balance [m w.e. yr^-1]', size=10)
#    axes[0,0].set_ylim(-210, 50)
#    axes[1,0].set_ylim(-8.5, 1.5)
    axes[0,0].set_ylim(-400, 200)
    axes[1,0].set_ylim(-12, 5)
    
    # tick marks for m w.e.    
    axes[1,0].yaxis.set_major_locator(plt.MultipleLocator(2))
    axes[1,0].yaxis.set_minor_locator(plt.MultipleLocator(0.5))
    
    # tick marks for Gt
    axes[0,0].yaxis.set_major_locator(plt.MultipleLocator(40))
    axes[0,0].yaxis.set_minor_locator(plt.MultipleLocator(10))
    
    # save figure
    fig.set_size_inches(5, 6)
    figure_fp = os.getcwd() + '/../Output/plots/gcm_compare_parameters/'
    if os.path.exists(figure_fp) == False:
        os.makedirs(figure_fp)
    figure_fn = 'massbalparams_tidewater_plot_' + rcp_list[j] + '.png'
    fig.savefig(figure_fp + figure_fn, bbox_inches='tight', dpi=300)
        
        
