# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Imports

import numpy as np
import xarray as xr
import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.patches import Polygon
from matplotlib import colors as mat_colors
import mpl_toolkits.axisartist as axisartist
from mpl_toolkits.axes_grid1 import Size, Divider

# # Define colors

colors = sns.color_palette("colorblind")
colors

axis_color = list(colors[7]) + [1.]
color_1 = list(colors[3]) + [1.]
color_2 = list(colors[0]) + [1.]
color_3 = list(colors[4]) + [1.]
color_4 = list(colors[2]) + [1.]
glacier_color = glacier_color = list(colors[9]) + [.5]


# # Define functions

# ## Performance measurements

# +
def BIAS(a1, a2):
    return (a1 - a2).mean().item()


def RMSE(a1, a2):
    return np.sqrt(((a1 - a2)**2).mean()).item()


def DIFF(a1, a2):
    return np.max(np.abs(a1 - a2)).item()


# -

# ## heatmap

def heatmap(datasets, # first_dataset, second_dataset,
            opti_var,
            annotation=None,
            annotation_x_position=0,
            fig=None, ax=None,
            cmap='vlag',
            cmap_levels=None,
            grid_color='grey',
            grid_linewidth=1.5,
            presentation=False,
            labels_pad=-360,
            xlim=None,
            nr_of_iterations=None):
    if not ax:
        ax = plt.gca()
    
    if not fig:
        fig = plt.gcf()

    if all(dataset is None for dataset in datasets):
        raise ValueError('All datasets are None!')

    # define variables for plotting
    guess_opti_var = []
    first_guess_diff = []
    true_opti_var = []
    BIAS_opti_var = []
    RMSE_opti_var = []
    DIFF_opti_var = []
    fct_opti_var = []
    times = []
    maxiters = []
    BIAS_sfc = []
    RMSE_sfc = []
    DIFF_sfc = []
    BIAS_w = []
    RMSE_w = []
    DIFF_w = []
    BIAS_fg = []
    RMSE_fg = []
    DIFF_fg = []
    BIAS_sfc_fg = []
    RMSE_sfc_fg = []
    DIFF_sfc_fg = []
    array_length = 0
    check_first_guess = None
    check_true_opti_var = None

    # create data and label variables
    for dataset in datasets:
        # check if the current dataset contains data or if the data was not available
        if dataset is None:
            guess_opti_var.append(None)
            first_guess_diff.append(None)
            true_opti_var.append(None)
            BIAS_opti_var.append(None)
            RMSE_opti_var.append(None)
            DIFF_opti_var.append(None)
            fct_opti_var.append(None)
            times.append(None)
            maxiters.append(None)
            BIAS_sfc.append(None)
            RMSE_sfc.append(None)
            DIFF_sfc.append(None)
            BIAS_w.append(None)
            RMSE_w.append(None)
            DIFF_w.append(None)
        elif type(dataset) != xr.core.dataset.Dataset:  # if no minimisation possible
            guess_opti_var.append('no_minimisation')
            first_guess_diff.append(None)
            true_opti_var.append(None)
            BIAS_opti_var.append(None)
            RMSE_opti_var.append(None)
            DIFF_opti_var.append(None)
            fct_opti_var.append(None)
            times.append(None)
            maxiters.append(None)
            BIAS_sfc.append(None)
            RMSE_sfc.append(None)
            DIFF_sfc.append(None)
            BIAS_w.append(None)
            RMSE_w.append(None)
            DIFF_w.append(None)
        else:
            # find index corresponding to max time
            max_index = len(dataset['computing_time'].values) - 1

            if nr_of_iterations is not None:
                max_index = nr_of_iterations - 1
            elif xlim is not None:
                # only consider as many points until max of xlim is reached
                if dataset['computing_time'].values[-1] > xlim[1]:
                    max_index = np.argmax(dataset['computing_time'].values > xlim[1])

            if opti_var == 'bed_h':
                guess_opti_var.append((dataset.guessed_bed_h[max_index] - dataset.true_bed_h).values)
                first_guess_diff.append((dataset.first_guessed_bed_h - dataset.true_bed_h).values)
                true_opti_var.append(dataset.true_bed_h.values)
                BIAS_opti_var.append(BIAS(dataset.guessed_bed_h[max_index], dataset.true_bed_h))
                RMSE_opti_var.append(RMSE(dataset.guessed_bed_h[max_index], dataset.true_bed_h))
                DIFF_opti_var.append(DIFF(dataset.guessed_bed_h[max_index], dataset.true_bed_h))
                if check_first_guess is None:
                    BIAS_fg = BIAS(dataset.first_guessed_bed_h, dataset.true_bed_h)
                    RMSE_fg = RMSE(dataset.first_guessed_bed_h, dataset.true_bed_h)
                    DIFF_fg = DIFF(dataset.first_guessed_bed_h, dataset.true_bed_h)
            elif opti_var == 'bed_shape':
                guess_opti_var.append((dataset.guessed_bed_shape[-1] - dataset.true_bed_shape).values)
                first_guess_diff.append((dataset.first_guessed_bed_shape - dataset.true_bed_shape).values)
                true_opti_var.append(dataset.true_bed_shape.values)
                BIAS_opti_var.append(BIAS(dataset.guessed_bed_shape[max_index], dataset.true_bed_shape))
                RMSE_opti_var.append(RMSE(dataset.guessed_bed_shape[max_index], dataset.true_bed_shape))
                DIFF_opti_var.append(DIFF(dataset.guessed_bed_shape[max_index], dataset.true_bed_shape))
                if check_first_guess is None:
                    BIAS_fg = BIAS(dataset.first_guessed_bed_shape, dataset.true_bed_shape)
                    RMSE_fg = RMSE(dataset.first_guessed_bed_shape, dataset.true_bed_shape)
                    DIFF_fg = DIFF(dataset.first_guessed_bed_shape, dataset.true_bed_shape)
            elif opti_var == 'w0':
                guess_opti_var.append((dataset.guessed_w0[-1] - dataset.true_w0).values)
                first_guess_diff.append((dataset.first_guessed_w0 - dataset.true_w0).values)
                true_opti_var.append(dataset.true_w0.values)
                BIAS_opti_var.append(BIAS(dataset.guessed_w0[max_index], dataset.true_w0))
                RMSE_opti_var.append(RMSE(dataset.guessed_w0[max_index], dataset.true_w0))
                DIFF_opti_var.append(DIFF(dataset.guessed_w0[max_index], dataset.true_w0))
                if check_first_guess is None:
                    BIAS_fg = BIAS(dataset.first_guessed_w0, dataset.true_w0)
                    RMSE_fg = RMSE(dataset.first_guessed_w0, dataset.true_w0)
                    DIFF_fg = DIFF(dataset.first_guessed_w0, dataset.true_w0)
            else:
                raise ValueError('Unknown opti var!')
                
            fct_opti_var.append(dataset.function_calls[max_index].values)
            times.append(dataset.computing_time[max_index].values)
            maxiters.append(dataset.attrs['maxiter_reached'])
            BIAS_sfc.append(BIAS(dataset.surface_h[max_index], dataset.true_surface_h))
            RMSE_sfc.append(RMSE(dataset.surface_h[max_index], dataset.true_surface_h))
            DIFF_sfc.append(DIFF(dataset.surface_h[max_index], dataset.true_surface_h))
            BIAS_w.append(BIAS(dataset.widths[max_index], dataset.true_widths))
            RMSE_w.append(RMSE(dataset.widths[max_index], dataset.true_widths))
            DIFF_w.append(DIFF(dataset.widths[max_index], dataset.true_widths))

            # determine array length for empty lines
            if array_length == 0:
                array_length = dataset.points_with_ice[-1].values + 1
            # check that the arrays have the same number of points with ice
            elif array_length != dataset.points_with_ice[-1].values + 1:
                raise ValueError('Not the same lentgth of points with ice!!!')

            # check if all experiments start with the same true values and first guess
            # in the first round save values
            if check_first_guess is None:
                check_first_guess = first_guess_diff[-1]
                check_true_opti_var = true_opti_var[-1]
                
                # not implemented yet
                BIAS_sfc_fg = BIAS(dataset.first_guess_surface_h, dataset.true_surface_h)
                RMSE_sfc_fg = RMSE(dataset.first_guess_surface_h, dataset.true_surface_h)
                DIFF_sfc_fg = DIFF(dataset.first_guess_surface_h, dataset.true_surface_h)
                BIAS_w_fg = BIAS(dataset.first_guess_widths, dataset.true_widths)
                RMSE_w_fg = RMSE(dataset.first_guess_widths, dataset.true_widths)
                DIFF_w_fg = DIFF(dataset.first_guess_widths, dataset.true_widths)

            # after first round compare all values to first ones to make sure comparing the same start conditions
            else:
                if np.sum(check_true_opti_var - true_opti_var[-1]) != 0:
                    raise ValueError('Not the same true control variable!!!')
                if np.sum(check_first_guess - first_guess_diff[-1]) != 0:
                    raise ValueError('Not the same first guess!!!')

    # create variables for ploting (data and y label)
    data = []
    y_labels = []

    # first add heading
    data.append(np.empty((array_length)) * np.nan)
    if not presentation:
        if opti_var == 'bed_h':
            y_labels.append(r'    RMSE_b  DIFF_b  RMSE_s  DIFF_s  fct  $T_{CPU}$')
        elif opti_var in ['bed_shape', 'w0']:
            y_labels.append('    RMSE_b  DIFF_b  RMSE_w  DIFF_w  fct  time')
        else:
            raise ValueError('Unknown opti_var !')
        y_label_variable_format = '{:6.2f}, {: 6.2f}, {:6.2f}, {:6.2f}'
    else:
        if opti_var == 'bed_h':
            y_labels.append('     DIFF  DIFF_s  fct  time')
        elif opti_var in ['bed_shape', 'w0']:
            y_labels.append('     DIFF  DIFF_w  fct  time')
        else:
            raise ValueError('Unknown opti_var !')
        y_label_variable_format = '{: 6.2f}, {:6.2f}'

    if not presentation:
        # add first guess
        data.append(check_first_guess)
        if opti_var == 'bed_h':
            y_labels.append(('fg:' + y_label_variable_format).format(RMSE_fg, DIFF_fg,
                                                                     RMSE_sfc_fg, DIFF_sfc_fg))
        elif opti_var in ['bed_shape', 'w0']:
            y_labels.append(('fg:' + y_label_variable_format).format(RMSE_fg, DIFF_fg,
                                                                     RMSE_w_fg, DIFF_w_fg))
        else:
            raise ValueError('Unknown opti_var !')
    else:
        # add first guess
        data.append(check_first_guess)
        if opti_var == 'bed_h':
            y_labels.append(('fg:' + y_label_variable_format).format(DIFF_fg, DIFF_sfc_fg))
        elif opti_var in ['bed_shape', 'w0']:
            y_labels.append(('fg:' + y_label_variable_format).format(DIFF_fg, DIFF_w_fg))
        else:
            raise ValueError('Unknown opti_var !')
    
    # add two format placeholders for fct_calls and time
    y_label_variable_format += ', {:3d}, {:3.0f}s'

    # add all other data with empty line for None datasets
    for i, guess in enumerate(guess_opti_var):
        if guess is None:
            data.append(np.empty((array_length)) * np.nan)
            if i < 9:
                y_labels.append((' ' + str(i+1) + ':    NO DATAFILE FOUND'))
            else:
                y_labels.append((str(i+1) + ':    NO DATAFILE FOUND'))
        elif type(guess) is str:
            data.append(np.empty((array_length)) * np.nan)
            if i < 9:
                y_labels.append((' ' + str(i+1) + ':    NO Minimisation Possible'))
            else:
                y_labels.append((str(i+1) + ':    NO Minimisation Possible'))
        else:
            data.append(guess)
            if i < 9:
                y_label_text = (' ' + str(i+1) + ':' + y_label_variable_format)
            else:
                y_label_text = (str(i+1) + ':' + y_label_variable_format)
            
            if maxiters[i] == 'yes':
                y_label_text += '+'

            if opti_var == 'bed_h':
                if not presentation:
                    y_labels.append(y_label_text.format(RMSE_opti_var[i],
                                                        DIFF_opti_var[i],
                                                        RMSE_sfc[i],
                                                        DIFF_sfc[i],
                                                        fct_opti_var[i],
                                                        times[i]))
                else:
                    y_labels.append(y_label_text.format(DIFF_opti_var[i],
                                                        DIFF_sfc[i],
                                                        fct_opti_var[i],
                                                        times[i]))
            elif opti_var in ['bed_shape', 'w0']:
                if not presentation:
                    y_labels.append(y_label_text.format(RMSE_opti_var[i],
                                                        DIFF_opti_var[i],
                                                        RMSE_w[i],
                                                        DIFF_w[i],
                                                        fct_opti_var[i],
                                                        times[i]))
                else:
                    y_labels.append(y_label_text.format(DIFF_opti_var[i],
                                                        DIFF_w[i],
                                                        fct_opti_var[i],
                                                        times[i]))
            else:
                raise ValueError('Unknown opti_var !')

    # make data an numpy array
    data = np.array(data)

    #choose colormap
    if not cmap_levels:
        color_nr = 100
        if opti_var == 'bed_h':
            cmap_limit = np.max(np.abs(check_first_guess))
            #cmap_limit = np.max(np.array([np.abs(np.floor(np.nanmin(np.array(data)))),
            #                              np.abs(np.ceil(np.nanmax(np.array(data))))]))
        elif opti_var in ['bed_shape', 'w0']:
            cmap_limit = np.max(np.abs(check_first_guess))
            #cmap_limit = np.max(np.array([np.abs(np.floor(np.nanmin(np.array(data)) * 10)),
            #                              np.abs(np.ceil(np.nanmax(np.array(data)) * 10))])) / 10
        else:
            raise ValueError('Unknown opti var!!')
        #if (np.min(data) < 0) & (np.max(data) > 0):
        cmap_levels = np.linspace(-cmap_limit, cmap_limit, color_nr, endpoint=True)
        #elif (np.min(data) < 0) & (np.max(data) =< 0):
        #    cmap_levels = np.linspace(-cmap_limit, 0, color_nr, endpoint=True)
        #elif (np.min(data) >= 0) & (np.max(data) > 0)
    else:
        color_nr = len(cmap_levels) - 1
    
    rel_color_steps = np.arange(color_nr)/color_nr
    if cmap == 'rainbow':
        colors = cm.rainbow(rel_color_steps)
    elif cmap == 'vlag':
        colors = sns.color_palette('vlag', color_nr)
    elif cmap == 'icefire':
        colors = sns.color_palette('icefire', color_nr)
    elif cmap == 'Spectral':
        colors = sns.color_palette('Spectral_r', color_nr)
    
    cmap = LinearSegmentedColormap.from_list('custom', colors, N=color_nr)
    cmap.set_bad(color='white')
    norm = mat_colors.BoundaryNorm(cmap_levels, cmap.N)

    # plot heatmap
    im = plt.imshow(data, aspect='auto', interpolation=None, cmap=cmap, norm=norm, alpha=1.)
    
    # Turn spines and ticks off and create white frame.
    for key, spine in ax.axis.items():
        spine.major_ticks.set_visible(False)
        spine.minor_ticks.set_visible(False)
        spine.line.set_visible(False)
        # spine.line.set_color(grid_color)
        # spine.line.set_linewidth(0) #grid_linewidth)
    
    # set y ticks
    ax.set_yticks(np.arange(data.shape[0]))
                
    ax.set_yticklabels(y_labels)
    #for tick in ax.get_yticklabels():
    #    tick.set_fontname("Arial")

    # align yticklabels left
    ax.axis["left"].major_ticklabels.set_ha("left")
    
    # set pad to put labels over heatmap
    ax.axis["left"].major_ticklabels.set_pad(labels_pad)

    # set y minor grid
    ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)
    ax.grid(which="minor", axis='y', color=grid_color, linestyle='-', linewidth=grid_linewidth)
    
    # set x ticklabels off
    ax.set_xticklabels([])
    
    # create colorbar
    cax = ax.inset_axes([1.01, 0.1, 0.03, 0.8]) 
    #cax = fig.add_axes([0.5, 0, 0.01, 1])
    cbar = fig.colorbar(im, cax=cax, boundaries=cmap_levels, spacing='proportional',)
    cbar.set_ticks([np.min(cmap_levels),0,np.max(cmap_levels)])
    if opti_var == 'bed_h':
        cbar.set_ticklabels(['{:d}'.format(int(-cmap_limit)), '0' ,'{:d}'.format(int(cmap_limit))])
    elif opti_var == 'bed_shape':
        cbar.set_ticklabels(['{:.1f}'.format(-cmap_limit), '0' ,'{:.1f}'.format(cmap_limit)])
    elif opti_var == 'w0':
        cbar.set_ticklabels(['{:d}'.format(int(-cmap_limit)), '0' ,'{:d}'.format(int(cmap_limit))])
    else:
        raise ValueError('Unknown opti var!!')
    #cbar.ax.set_ylabel(cbarlabel,)
    
    # set title
    #ax.set_title(title)
    
    if annotation is not None:
        # include text
        ax.text(annotation_x_position, 1, 
                 annotation,
                 horizontalalignment='center',
                 verticalalignment='center',
                 transform=ax.transAxes)
    
    return im


# ## performance plot

def performance_plot(axs,
                     datasets,
                     annotation=None,
                     fig=None,
                     # 'bed_h RMSE', 'bed_h Diff', 'bed_h Bias',
                     # 'bed_shape RMSE', 'bed_shape Diff', 'bed_shape Bias',
                     # 'w0 RMSE', 'w0 Diff', 'w0 Bias',
                     # 'sfc_h RMSE', 'sfc_h Diff', 'sfc_h Bias',
                     # 'widths RMSE', 'widths Diff', 'widths Bias'
                     performance_measurements=['bed_h RMSE', 'sfc_h RMSE'],
                     colors=['tab:blue', 'tab:red'],
                     xlim=[0,60],
                     line_width=2,
                     marker_size=5,
                     ylims=None,
                     nr_of_iterations=None
                    ):
    if not fig:
        fig = plt.gcf()

    all_x = []
    all_y = []

    nr_of_color = 0
    for nr_measure, measure in enumerate(performance_measurements):
        for dataset in datasets:
            if dataset is not None:
                max_index = len(dataset['computing_time'].values) - 1
                if nr_of_iterations is not None:
                    max_index = nr_of_iterations - 1
                elif xlim is not None:
                    # only consider as many points until max of xlim is reached
                    if dataset['computing_time'].values[-1] > xlim[1]:
                        max_index = np.argmax(dataset['computing_time'].values > xlim[1])

                # include time 0 for first guess
                tmp_x = [0]

                # add first guess values 
                if measure == 'bed_h RMSE':
                    tmp_y = [RMSE(dataset['first_guessed_bed_h'], dataset['true_bed_h'])]
                elif measure == 'bed_h Diff':
                    tmp_y = [DIFF(dataset['first_guessed_bed_h'], dataset['true_bed_h'])]
                elif measure == 'bed_h Bias':
                    tmp_y = [BIAS(dataset['first_guessed_bed_h'], dataset['true_bed_h'])]

                elif measure == 'bed_shape RMSE':
                    tmp_y = [RMSE(dataset['first_guessed_bed_shape'], dataset['true_bed_shape'])]
                elif measure == 'bed_shape Diff':
                    tmp_y = [DIFF(dataset['first_guessed_bed_shape'], dataset['true_bed_shape'])]
                elif measure == 'bed_shape Bias':
                    tmp_y = [BIAS(dataset['first_guessed_bed_shape'], dataset['true_bed_shape'])]

                elif measure == 'w0 RMSE':
                    tmp_y = [RMSE(dataset['first_guessed_w0'], dataset['true_w0'])]
                elif measure == 'w0 Diff':
                    tmp_y = [DIFF(dataset['first_guessed_w0'], dataset['true_w0'])]
                elif measure == 'w0 Bias':
                    tmp_y = [BIAS(dataset['first_guessed_w0'], dataset['true_w0'])]

                elif measure == 'sfc_h RMSE':
                    tmp_y = [RMSE(dataset['first_guess_surface_h'], dataset['true_surface_h'])]
                elif measure == 'sfc_h Diff':
                    tmp_y = [DIFF(dataset['first_guess_surface_h'], dataset['true_surface_h'])]
                elif measure == 'sfc_h Bias':
                    tmp_y = [DIFF(dataset['first_guess_surface_h'], dataset['true_surface_h'])]

                elif measure == 'widths RMSE':
                    tmp_y = [RMSE(dataset['first_guess_widths'], dataset['true_widths'])]
                elif measure == 'widths Diff':
                    tmp_y = [DIFF(dataset['first_guess_widths'], dataset['true_widths'])]
                elif measure == 'widths Bias':
                    tmp_y = [DIFF(dataset['first_guess_widths'], dataset['true_widths'])]

                else:
                    raise ValueError('Unknown performance measurement!')

                for i in dataset.coords['nr_of_iteration'].values[:max_index + 1] - 1:
                    tmp_x.append(dataset['computing_time'][i])
                    if measure == 'bed_h RMSE':
                        tmp_y.append(RMSE(dataset['guessed_bed_h'][i], dataset['true_bed_h']))
                        ylabel_1 = 'bed_h m'
                    elif measure == 'bed_h Diff':
                        tmp_y.append(DIFF(dataset['guessed_bed_h'][i], dataset['true_bed_h']))
                    elif measure == 'bed_h Bias':
                        tmp_y.append(BIAS(dataset['guessed_bed_h'][i], dataset['true_bed_h']))

                    elif measure == 'bed_shape RMSE':
                        tmp_y.append(RMSE(dataset['guessed_bed_shape'][i], dataset['true_bed_shape']))
                        ylabel_1 = 'bed_shape'
                    elif measure == 'bed_shape Diff':
                        tmp_y.append(DIFF(dataset['guessed_bed_shape'][i], dataset['true_bed_shape']))
                    elif measure == 'bed_shape Bias':
                        tmp_y.append(BIAS(dataset['guessed_bed_shape'][i], dataset['true_bed_shape']))

                    elif measure == 'w0 RMSE':
                        tmp_y.append(RMSE(dataset['guessed_w0'][i], dataset['true_w0']))
                        ylabel_1 = 'w0 m'
                    elif measure == 'w0 Diff':
                        tmp_y.append(DIFF(dataset['guessed_w0'][i], dataset['true_w0']))
                    elif measure == 'w0 Bias':
                        tmp_y.append(BIAS(dataset['guessed_w0'][i], dataset['true_w0']))

                    elif measure == 'sfc_h RMSE':
                        tmp_y.append(RMSE(dataset['surface_h'][i], dataset['true_surface_h']))
                        ylabel_2 = 'sfc_h m'
                    elif measure == 'sfc_h Diff':
                        tmp_y.append(DIFF(dataset['surface_h'][i], dataset['true_surface_h']))
                    elif measure == 'sfc_h Bias':
                        tmp_y.append(BIAS(dataset['surface_h'][i], dataset['true_surface_h']))

                    elif measure == 'widths RMSE':
                        tmp_y.append(RMSE(dataset['widths'][i], dataset['true_widths']))
                        ylabel_2 = 'widths m'
                    elif measure == 'widths Diff':
                        tmp_y.append(DIFF(dataset['widths'][i], dataset['true_widths']))
                    elif measure == 'widths Bias':
                        tmp_y.append(BIAS(dataset['widths'][i], dataset['true_widths']))

                    else:
                        raise ValueError('Unknown performance measurement!')
            else:
                tmp_x = []
                tmp_y = []

            axs[nr_measure].plot(tmp_x, tmp_y, '.-',
                                 color=colors[nr_of_color],
                                 lw=line_width,
                                 ms=marker_size)
            nr_of_color += 1

    axs[0].set_xlabel('time in s')
    axs[0].set_ylabel(ylabel_1, color=colors[0])
    axs[0].tick_params(axis='y', labelcolor=colors[0])
    axs[0].grid()

    if len(colors) == 2:
        axs[1].set_ylabel(ylabel_2, color=colors[1])
        axs[1].tick_params(axis='y', labelcolor=colors[1])
    elif len(colors) == 4:
        axs[1].set_ylabel(ylabel_2, color=colors[2])
        axs[1].tick_params(axis='y', labelcolor=colors[2])
    
    #for x, y, description in zip(all_x, all_y, descriptions):
    #    ax.plot(x, y, '.-', label=description)

    # ax.legend((),(),title=measure, loc='best')
    # ax.axvline(60, alpha=0.5, c='gray', ls='--')
    # ax.axvline(20, alpha=0.5, c='gray', ls='--')
    if xlim is not None:
        axs[0].set_xlim(xlim)
    
    if ylims is not None:
        axs[0].set_ylim(ylims[0])
        axs[0].set_yticks(np.linspace(ylims[0][0], ylims[0][1], 3))
        axs[1].set_ylim(ylims[1])
        axs[1].set_yticks(np.linspace(ylims[1][0], ylims[1][1], 3))


# ## performance plot cost function

def performance_cost_fkt(ax,
                         dataset,
                         nr_of_iterations,
                         annotation=None,
                         annotation_x_position=-0.2,
                         lw=2,
                         fontsize=25,
                         ms=10):
    x_use = np.arange(0, nr_of_iterations + 1)
    
    # calculate cost for first guess
    cost_term_1_fg = dataset.attrs['reg_parameters'][0] * np.power(
        dataset.true_surface_h.data[dataset.ice_mask.data] -
        dataset.first_guess_surface_h.data[dataset.ice_mask.data], 2).sum()
    cost_term_2_fg = dataset.attrs['reg_parameters'][0] * np.power(
        dataset.true_surface_h.data[~dataset.ice_mask.data] -
        dataset.first_guess_surface_h.data[~dataset.ice_mask.data], 2).sum()
    total_first_guess_bed_h = np.append(dataset.first_guessed_bed_h.data,
                                        dataset.total_true_bed_h.data[~dataset.ice_mask.data])
    db_dx = (total_first_guess_bed_h[1:] -
             total_first_guess_bed_h[:-1]) / dataset.attrs['distance_between_points_m']
    cost_term_3_fg = dataset.attrs['reg_parameters'][3] * np.power(db_dx, 2).sum()

    # put first guess cost together with rest
    cost_term_1 = np.append(cost_term_1_fg, dataset.cost_terms[:nr_of_iterations, 0].data)
    cost_term_2 = np.append(cost_term_2_fg, dataset.cost_terms[:nr_of_iterations, 2].data)
    cost_term_3 = np.append(cost_term_3_fg, dataset.cost_terms[:nr_of_iterations, 3].data)
    cost = np.append(cost_term_1_fg + cost_term_3_fg, 
                     dataset.cost[:nr_of_iterations].data)
    
    # plot curves
    ax.plot(x_use, cost,
            '.-',
            c=color_2,
            lw=lw,
            ms=ms,
            label='total cost',
            zorder=4)
    ax.plot(x_use, cost_term_1,
            '.-',
            c=color_1,
            lw=lw,
            ms=ms,
            label='surface misfit',
            zorder=1)
    ax.plot(x_use, cost_term_2,
            '.-',
            c=color_4,
            lw=lw,
            ms=ms,
            label='ice outside',
            zorder=2)
    ax.plot(x_use, cost_term_3,
            '.-',
            c=color_3,
            lw=lw,
            ms=ms,
            label='smoothing',
            zorder=3)
    
    ax.set_xlabel('Iteration', fontsize=fontsize, c=axis_color)
    ax.set_ylabel('cost value', fontsize=fontsize, c=axis_color)
    
    #[t.set_color(axis_color) for t in ax.xaxis.get_ticklines()]
    #[t.set_color(axis_color) for t in ax.xaxis.get_ticklabels()]
    ax.tick_params(axis='both', colors=axis_color, width=lw)
    ax.spines['bottom'].set_color(axis_color)
    ax.spines['bottom'].set_linewidth(lw)
    ax.spines['left'].set_color(axis_color)
    ax.spines['left'].set_linewidth(lw)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.legend(fontsize=fontsize)
    
    if annotation is not None:
        ax.text(annotation_x_position, 1,
                annotation,
                horizontalalignment='center',
                verticalalignment='center',
                transform = ax.transAxes)


# ## help function for heatmap axis

def setup_axes(fig, rect):
    ax = axisartist.Subplot(fig, rect)
    fig.add_subplot(ax)

    return ax


# ## convert iterations to datasets

def convert_iterations_to_dataset(dataset, timesteps=1):
    #total_iterations = dataset.coords['nr_of_iteration'].values
    #relative_iteration = int(total_iterations[-1] / (timesteps - 1))
    #plot_iterations = total_iterations[::relative_iteration]
    #if len(plot_iterations) == timesteps - 1:
    #    plot_iterations= np.append(plot_iterations, total_iterations[-1])
    #elif len(plot_iterations) == timesteps:
    #    plot_iterations[-1] = total_iterations[-1]
    #else:
    #    raise ValueError('Something went wrong with calculation of timesteps of iterations!')
    plot_iterations = dataset.coords['nr_of_iteration'].values

    timeseries_datasets = []
    for iteration in plot_iterations - 1:
        tmp_dataset = xr.Dataset(
                    data_vars={
                        'true_bed_h':
                            (['points_with_ice'],
                             dataset.true_bed_h),
                        'first_guessed_bed_h':
                            (['points_with_ice'],
                             dataset.first_guessed_bed_h),
                        'guessed_bed_h':
                            (['nr_of_iteration', 'points_with_ice'],
                             np.array([dataset.guessed_bed_h[iteration].values])),
                        'function_calls':
                            (['nr_of_iteration'],
                             np.array([dataset.function_calls[iteration].values])),
                        'true_surface_h':
                            (['total_distance'],
                             dataset.true_surface_h),
                        'first_guess_surface_h':
                            (['total_distance'],
                             dataset.first_guess_surface_h),
                        'surface_h':
                            (['nr_of_iteration', 'total_distance'],
                             np.array([dataset.surface_h[iteration].values])),
                        'true_widths':
                            (['total_distance'],
                             dataset.true_widths),
                        'first_guess_widths':
                            (['total_distance'],
                             dataset.first_guess_widths),
                        'widths':
                            (['nr_of_iteration', 'total_distance'],
                             np.array([dataset.widths[iteration].values])),
                        'computing_time':
                            (['nr_of_iteration'],
                             np.array([dataset.computing_time[iteration].values]))
                    },
                    coords={
                        'total_distance': dataset.coords['total_distance'],
                        'points_with_ice': dataset.coords['points_with_ice'],
                        'nr_of_iteration': ([1])
                    },
                    attrs={
                        #'computing_time': -1,
                        'maxiter_reached': False
                    }
        )
        #if iteration == plot_iterations[-1] - 1:
        #    tmp_dataset.attrs['computing_time'] = dataset.attrs['computing_time']
        timeseries_datasets.append(tmp_dataset)
    return timeseries_datasets


# ## function for profil

def profil_subplot(
            ax,
            true_bed_h,
            true_sfc_h,
            second_bed_h,
            second_sfc_h,
            title,
            fontsize,
            tick_labels=None,
            tick_labels_indexes=None,
            tick_labels_position='above',
            len_tick_label_line = 110,
            tick_label_text_y_distance = 2,
            extra_distance_between_tick_labels = 15,
            index_start=0,
            index_end=10,
            lw=2,
            ms=20,
            add_legend=False):
    # define index which points should be used
    index = np.arange(index_start, index_end)
    
    # plot true bed_h and surface_h
    ax.plot(total_distance[index],
            true_bed_h[index],
            '.-',
            lw=lw,
            ms=ms,
            c=color_1,
            label=r'true $b$')
    ax.plot(total_distance[index],
            true_sfc_h[index],
            '.--',
            lw=lw,
            ms=ms,
            c=color_1,
            label=r'$s^{e}_o$')
    
    # plot second bed_h and surface_h
    ax.plot(total_distance[index],
            np.append(second_bed_h, true_bed_h[~ice_mask])[index],
            '.-',
            lw=lw,
            ms=ms,
            c=color_2,
            zorder=5,
            label=r'guessed $b$')
    ax.plot(total_distance[index],
            second_sfc_h[index],
            '.--',
            lw=lw,
            ms=ms,
            c=color_2,
            label=r'$s^{e}_m$')
    
    # add glacier polygon
    x_use=total_distance[index]
    x_polygon = np.concatenate((x_use, x_use[::-1]))
    y_polygon = np.concatenate((np.append(second_bed_h, true_bed_h[~ice_mask])[index], 
                                second_sfc_h[index][::-1]))
    coord_polygon = np.concatenate((np.expand_dims(x_polygon, axis=1),np.expand_dims(y_polygon, axis=1)), axis=1)
    ax.add_patch(Polygon(coord_polygon,
                         fc=glacier_color,  
                         ec=None,#outline_color,
                         closed=False,
                         lw = 0.8,
                         zorder=1,
                         label=''))

    # add labels for grid points
    for i, (tick_label, point_index)in enumerate(zip(tick_labels, tick_labels_indexes)):
        x_use = total_distance[point_index]
        if tick_labels_position == 'above':
            y_use_upper = true_bed_h[point_index] + len_tick_label_line + i * extra_distance_between_tick_labels
            y_use_lower = np.min([true_bed_h[point_index], second_bed_h[point_index]])
            y_use_label = (true_bed_h[point_index] + len_tick_label_line + tick_label_text_y_distance +
                           i * extra_distance_between_tick_labels)
        elif tick_labels_position == 'below':
            y_use_upper = np.max([true_sfc_h[point_index], second_sfc_h[point_index]])
            y_use_lower = true_bed_h[point_index] - len_tick_label_line - i * extra_distance_between_tick_labels
            y_use_label = (true_bed_h[point_index] - len_tick_label_line - tick_label_text_y_distance -
                           i * extra_distance_between_tick_labels)
        # plot line
        ax.plot([x_use, x_use],
                [y_use_lower, y_use_upper],
               '-',
               c=axis_color,
               lw=lw-1,
               zorder=1)
        
        # add tick label text
        ax.text(x_use,
                y_use_label,
                tick_label,
                fontsize=fontsize,
                c=axis_color,
                verticalalignment='center',
                horizontalalignment='center')
    
    # add text with description
    ax.text(0.95, 0.95,
            title,
            fontsize=fontsize,
            verticalalignment='top',
            horizontalalignment='right',
            transform = ax.transAxes)
    # add visual axis
    x_origin = total_distance[index[0]]
    z_origin = true_bed_h[index[-1]]
    x_len = 0.12
    z_len = 50
    x_text_setoff = 0.015
    z_text_setoff = 10

    # add z axis
    plt.annotate(text='',
                 xy=(x_origin, z_origin), 
                 xytext=(x_origin, z_origin + z_len),
                 arrowprops=dict(arrowstyle='<-',
                                 mutation_scale=25,
                                 color=axis_color,
                                 lw=1),
                 zorder=1
                 )
    plt.text(x_origin ,z_origin + z_len + z_text_setoff,'z',
             horizontalalignment='center',
             verticalalignment='center',
             fontsize=fontsize + 2,
             c=axis_color)

    # add x axis
    plt.annotate(text='',
                 xy=(x_origin, z_origin), 
                 xytext=(x_origin + x_len, z_origin),
                 arrowprops=dict(arrowstyle='<-',
                                 mutation_scale=25,
                                 color=axis_color,
                                 lw=1),
                 zorder=1
                 )
    plt.text(x_origin + x_len + x_text_setoff, z_origin, 'x',
             horizontalalignment='center',
             verticalalignment='center',
             fontsize=fontsize + 2,
             c=axis_color)
    # set limits of x and y axis
    #ax.set_xlim([5.5,6.6])
    #ax.set_ylim([2340, 2720])
    
    if add_legend:
        ax.legend(fontsize=fontsize)
    ax.set_xticks([])
    ax.set_yticks([])


# ## profil legend

def add_legend(ax,
               title,
               fontsize,
               lw,
               ms):
    # plot true bed_h and surface_h
    ax.plot([],
            [],
            '.-',
            lw=lw,
            ms=ms,
            c=color_1,
            label=r'$b_{t}$')
    ax.plot([],
            [],
            '.--',
            lw=lw,
            ms=ms,
            c=color_1,
            label=r'$s^{e}_{o}$')
    
    # plot second bed_h and surface_h
    ax.plot([],
            [],
            '.-',
            lw=lw,
            ms=ms,
            c=color_2,
            zorder=5,
            label=r'$b$')
    ax.plot([],
            [],
            '.--',
            lw=lw,
            ms=ms,
            c=color_2,
            label=r'$s^{e}_{m}$')
    
    ax.legend(loc='center', fontsize=fontsize, title=title)
    ax.axis('off')


# # Define figure function

def create_figure(dataset,
                  nr_of_iterations = 10,
                  facecolor = 'white',
                  labels_pad = -620,
                  cmap = 'Spectral',
                  fontsize = 25,
                  lw=2,
                  ms=20,
                  annotation_x_position_spatial = -0.1,
                  annotation_x_position_performance = -0.2,
                  index_start_first_profil_row = 0,
                  index_end_first_profil_row = 6,
                  index_start_second_profil_row = 65,
                  index_end_second_profil_row = 71,
                  save_file = False,
                  filename = 'iterative_behaviour_rec_lin_cons_ret.pdf'):

    #plt.rcParams['font.family'] = 'monospace'
    mpl.rcParams.update({'font.size': fontsize})

    fig = plt.figure(figsize=(1,1), facecolor='white')

    # define grid
    # define fixed size of subplot
    spatial_performance_height = 5
    spatial_performance_separation_x = 3
    spatial_profil_separation_y = 1

    profil_subplot_width = 5
    profil_subplot_height = 5
    profil_subplot_separation_x = .1
    profil_subplot_separation_y = .3

    legend_width = 3

    # calculate grid spacing
    delta_different_x_separation = (spatial_performance_separation_x - profil_subplot_separation_x) / 2

    # fixed size in inch
    # along x axis                                                              x-index for locator
    horiz = [Size.Fixed(legend_width),                                          # 0 legend
             Size.Fixed(profil_subplot_width),                                  # 1 1st profil, spatial start
             Size.Fixed(profil_subplot_separation_x),   
             Size.Fixed(profil_subplot_width - delta_different_x_separation),   # 3 2nd profil start
             Size.Fixed(delta_different_x_separation),                          # 4 spatial end
             Size.Fixed(profil_subplot_separation_x),                           # 5 2nd profil end
             Size.Fixed(delta_different_x_separation),                          # 6 3rd profil start
             Size.Fixed(profil_subplot_width - delta_different_x_separation),   # 7 performance start
             Size.Fixed(profil_subplot_separation_x),                           # 8 3rd profil end
             Size.Fixed(profil_subplot_width),                                  # 9 4th profil, 
                                                                                # 10 performance end
            ]
                                                                              # y-index for locator
    vert = [Size.Fixed(profil_subplot_height),                                # 0 2nd row profil
            Size.Fixed(profil_subplot_separation_y),
            Size.Fixed(profil_subplot_height),                                # 2 1st row profil
            Size.Fixed(spatial_profil_separation_y),
            Size.Fixed(spatial_performance_height)                            # 4 spatial/performance row
           ]

    rect = (0., 0., 1., 1.)  # Position of the grid in the figure

    # divide the axes rectangle into grid whose size is specified by horiz * vert
    divider = Divider(fig, rect, horiz, vert, aspect=False)

    # add spatial plot
    plot_datasets = convert_iterations_to_dataset(dataset)

    with plt.rc_context({'font.family': 'monospace'}):
        ax = setup_axes(fig, 111)
        im = heatmap(plot_datasets[:nr_of_iterations],
                     opti_var='bed_h',
                     annotation='(a)',
                     annotation_x_position=annotation_x_position_spatial,
                     fig=fig,
                     ax=ax,
                     cmap=cmap,
                     grid_color=facecolor,
                     presentation=False,
                     labels_pad=labels_pad,
                     xlim=None,
                     nr_of_iterations=None)
        ax.set_axes_locator(divider.new_locator(nx=1, nx1=4, ny=4))

    # add performance plot for cost function terms
    ax = fig.subplots()
    performance_cost_fkt(ax,
                         dataset,
                         nr_of_iterations,
                         annotation='(b)',
                         annotation_x_position=annotation_x_position_performance,
                         lw=lw,
                         fontsize=fontsize,
                         ms=ms)
    ax.set_axes_locator(divider.new_locator(nx=7, nx1=10, ny=4))

    # get some data out of dataset
    true_bed_h = dataset.total_true_bed_h.data
    true_sfc_h = dataset.true_surface_h.data
    first_guess_bed_h = dataset.first_guessed_bed_h.data
    first_guess_sfc_h = dataset.first_guess_surface_h.data
    guessed_bed_h = dataset.guessed_bed_h.data
    guessed_sfc_h = dataset.surface_h.data
    total_distance = dataset.coords['total_distance'].data
    ice_mask = dataset.ice_mask.data

    # add first row profil plots
    tick_labels=[r'$1$', r'$2$']
    tick_labels_indexes=[0, 1]
    tick_labels_position='below'
    len_tick_label_line = 20
    tick_label_text_y_distance = 10
    extra_distance_between_tick_labels = 0

    # first guess
    ax = fig.subplots()
    profil_subplot(
        ax=ax,
        true_bed_h=true_bed_h,
        true_sfc_h=true_sfc_h,
        second_bed_h=first_guess_bed_h,
        second_sfc_h=first_guess_sfc_h,
        title='(c) first guess',
        fontsize=fontsize,
        tick_labels=tick_labels,
        tick_labels_indexes=tick_labels_indexes,
        tick_labels_position=tick_labels_position,
        len_tick_label_line=len_tick_label_line,
        tick_label_text_y_distance=tick_label_text_y_distance,
        extra_distance_between_tick_labels=extra_distance_between_tick_labels,
        index_start=index_start_first_profil_row,
        index_end=index_end_first_profil_row,
        lw=lw,
        ms=ms)
    ax.set_axes_locator(divider.new_locator(nx=1, ny=2))

    # add iteration 1, 2, 3
    for i, prefix, nx, nx1 in zip(np.arange(1,4), ['(d) ', '(e) ', '(f) '], [3, 6, 9], [5, 8, 10]):
        ax = fig.subplots()
        profil_subplot(
            ax=ax,
            true_bed_h=true_bed_h,
            true_sfc_h=true_sfc_h,
            second_bed_h=guessed_bed_h[i-1],
            second_sfc_h=guessed_sfc_h[i-1],
            title=prefix + str(i) + '. Iteration',
            fontsize=fontsize,
            tick_labels=tick_labels,
            tick_labels_indexes=tick_labels_indexes,
            tick_labels_position=tick_labels_position,
            len_tick_label_line=len_tick_label_line,
            tick_label_text_y_distance=tick_label_text_y_distance,
            extra_distance_between_tick_labels=extra_distance_between_tick_labels,
            index_start=index_start_first_profil_row,
            index_end=index_end_first_profil_row,
            lw=lw,
            ms=ms)
        ax.set_axes_locator(divider.new_locator(nx=nx, nx1=nx1, ny=2))

    # add second row profil plots
    tick_labels=[r'$trm$', r'$trm-1$']
    tick_labels_indexes=[68, 67]
    tick_labels_position='above'
    len_tick_label_line = 35
    tick_label_text_y_distance = 10
    extra_distance_between_tick_labels = 10
    # first guess
    ax = fig.subplots()
    profil_subplot(
        ax=ax,
        true_bed_h=true_bed_h,
        true_sfc_h=true_sfc_h,
        second_bed_h=first_guess_bed_h,
        second_sfc_h=first_guess_sfc_h,
        title='(g) first guess',
        fontsize=fontsize,
        tick_labels=tick_labels,
        tick_labels_indexes=tick_labels_indexes,
        tick_labels_position=tick_labels_position,
        len_tick_label_line=len_tick_label_line,
        tick_label_text_y_distance=tick_label_text_y_distance,
        extra_distance_between_tick_labels=extra_distance_between_tick_labels,
        index_start=index_start_second_profil_row,
        index_end=index_end_second_profil_row,
        lw=lw,
        ms=ms)
    ax.set_axes_locator(divider.new_locator(nx=1, ny=0))

    # add iteration 1, 2, 3
    for i, prefix, nx, nx1 in zip(np.arange(1,4), ['(h) ', '(i) ', '(j) '], [3, 6, 9], [5, 8, 10]):
        ax = fig.subplots()
        profil_subplot(
            ax=ax,
            true_bed_h=true_bed_h,
            true_sfc_h=true_sfc_h,
            second_bed_h=guessed_bed_h[i-1],
            second_sfc_h=guessed_sfc_h[i-1],
            title=prefix + str(i) + '. Iteration',
            fontsize=fontsize,
            tick_labels=tick_labels,
            tick_labels_indexes=tick_labels_indexes,
            tick_labels_position=tick_labels_position,
            len_tick_label_line=len_tick_label_line,
            tick_label_text_y_distance=tick_label_text_y_distance,
            extra_distance_between_tick_labels=extra_distance_between_tick_labels,
            index_start=index_start_second_profil_row,
            index_end=index_end_second_profil_row,
            lw=lw,
            ms=ms)
        ax.set_axes_locator(divider.new_locator(nx=nx, nx1=nx1, ny=0))

    # add profil legend
    ax = fig.subplots()
    add_legend(ax=ax,
               title='For (c) - (j)',
               fontsize=fontsize,
               lw=lw,
               ms=ms)
    # add text for first row
    ax.text(0, 0.75,
            'glacier top:',
            fontsize=fontsize,
            verticalalignment='center',
            horizontalalignment='left',
            transform = ax.transAxes)

    # add text for secont row
    ax.text(0, 0.25,
            'glacier terminus:',
            fontsize=fontsize,
            verticalalignment='center',
            horizontalalignment='left',
            transform = ax.transAxes)

    ax.set_axes_locator(divider.new_locator(nx=0, ny=0, ny1=3))

    if save_file:
        fig.savefig(filename, format='pdf', bbox_inches='tight', dpi=300);


# # Plot it

# +
input_folder = 'plot_data/'
filename_nc = 'rec_line_cons_ret_bed_hreg10.nc'

with xr.open_dataset(input_folder + filename_nc) as ds:
    dataset = ds

# define some variables from the dataset needed globaly by the plot functions (could be made nicer;))
true_bed_h = dataset.total_true_bed_h.data
true_sfc_h = dataset.true_surface_h.data
first_guess_bed_h = dataset.first_guessed_bed_h.data
first_guess_sfc_h = dataset.first_guess_surface_h.data
guessed_bed_h = dataset.guessed_bed_h.data
guessed_sfc_h = dataset.surface_h.data
total_distance = dataset.coords['total_distance'].data
ice_mask = dataset.ice_mask.data
    
create_figure(dataset,
              nr_of_iterations = 10,
              facecolor = 'white',
              labels_pad = -620,
              cmap = 'Spectral',
              fontsize = 25,
              lw=2,
              ms=20,
              annotation_x_position_spatial = -0.1,
              annotation_x_position_performance = -0.2,
              index_start_first_profil_row = 0,
              index_end_first_profil_row = 6,
              index_start_second_profil_row = 65,
              index_end_second_profil_row = 71,
              save_file = True,
              filename = 'iterative_behaviour_rec_lin_cons_ret.pdf')
# -



StarcoderdataJupyterScriptsDedupFiltered