import itertools
import os

import matplotlib.legend as mlegend
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Ellipse
from matplotlib.patches import Polygon
from matplotlib.patches import Rectangle
from mpl_toolkits.axes_grid1 import ImageGrid
from pypet import Trajectory
from pypet.brian2 import Brian2MonitorResult

from scripts.spatial_maps.spatial_network_layout import Interneuron, get_position_mesh
from scripts.spatial_maps.correlation_length_fit.correlation_length_fit import \
    correlation_length_fit_dict
from scripts.model_figure.figure_utils import remove_frame, remove_ticks, add_length_scale, cm_per_inch, \
    panel_size, head_direction_input_colormap
from scripts.spatial_network.perlin_map.run_simulation_perlin_map import DATA_FOLDER, TRAJ_NAME, \
    get_input_head_directions, POLARIZED, CIRCULAR, NO_SYNAPSES

plt.style.use('../../model_figure/figures.mplstyle')

FIGURE_SAVE_PATH = '../../../figures/' + TRAJ_NAME + '/'
FIGURE_SAVE_FORMAT = '.pdf'
save_figs = False


def tablelegend(ax, col_labels=None, row_labels=None, title_label="", *args, **kwargs):
    """
    Place a table legend on the axes.

    Creates a legend where the labels are not directly placed with the artists,
    but are used as row and column headers, looking like this:

    title_label   | col_labels[1] | col_labels[2] | col_labels[3]
    -------------------------------------------------------------
    row_labels[1] |
    row_labels[2] |              <artists go there>
    row_labels[3] |


    Parameters
    ----------

    ax : `matplotlib.axes.Axes`
        The artist that contains the legend table, i.e. current axes instant.

    col_labels : list of str, optional
        A list of labels to be used as column headers in the legend table.
        `len(col_labels)` needs to match `ncol`.

    row_labels : list of str, optional
        A list of labels to be used as row headers in the legend table.
        `len(row_labels)` needs to match `len(handles) // ncol`.

    title_label : str, optional
        Label for the top left corner in the legend table.

    ncol : int
        Number of columns.


    Other Parameters
    ----------------

    Refer to `matplotlib.legend.Legend` for other parameters.

    """
    #################### same as `matplotlib.axes.Axes.legend` #####################
    handles, labels, extra_args, kwargs = mlegend._parse_legend_args([ax], *args, **kwargs)
    if len(extra_args):
        raise TypeError('legend only accepts two non-keyword arguments')

    if col_labels is None and row_labels is None:
        ax.legend_ = mlegend.Legend(ax, handles, labels, **kwargs)
        ax.legend_._remove_method = ax._remove_legend
        return ax.legend_
    #################### modifications for table legend ############################
    else:
        ncol = kwargs.pop('ncol')
        handletextpad = kwargs.pop('handletextpad', 0 if col_labels is None else -2)
        title_label = [title_label]

        # blank rectangle handle
        extra = [Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor='none', linewidth=0)]

        # empty label
        empty = [""]

        # number of rows infered from number of handles and desired number of columns
        nrow = len(handles) // ncol

        # organise the list of handles and labels for table construction
        if col_labels is None:
            assert nrow == len(
                row_labels), "nrow = len(handles) // ncol = %s, but should be equal to len(row_labels) = %s." % (
                nrow, len(row_labels))
            leg_handles = extra * nrow
            leg_labels = row_labels
        elif row_labels is None:
            assert ncol == len(col_labels), "ncol = %s, but should be equal to len(col_labels) = %s." % (
                ncol, len(col_labels))
            leg_handles = []
            leg_labels = []
        else:
            assert nrow == len(
                row_labels), "nrow = len(handles) // ncol = %s, but should be equal to len(row_labels) = %s." % (
                nrow, len(row_labels))
            assert ncol == len(col_labels), "ncol = %s, but should be equal to len(col_labels) = %s." % (
                ncol, len(col_labels))
            leg_handles = extra + extra * nrow
            leg_labels = title_label + row_labels
        for col in range(ncol):
            if col_labels is not None:
                leg_handles += extra
                leg_labels += [col_labels[col]]
            leg_handles += handles[col * nrow:(col + 1) * nrow]
            leg_labels += empty * nrow

        # Create legend
        ax.legend_ = mlegend.Legend(ax, leg_handles, leg_labels, ncol=ncol + int(row_labels is not None),
                                    handletextpad=handletextpad, **kwargs)
        ax.legend_._remove_method = ax._remove_legend
        return ax.legend_


def get_closest_scale(traj, scale):
    available_lengths = sorted(list(set(traj.f_get("scale").f_get_range())))
    closest_length = available_lengths[np.argmin(np.abs(np.array(available_lengths) - scale))]
    if closest_length != scale:
        print("Warning: desired correlation length {:.1f} not available. Taking {:.1f} instead".format(
            scale, closest_length))
    corr_len = closest_length
    return corr_len


def get_closest_fit_correlation_length(traj, fit_corr_len):
    corr_len_fit_dict = correlation_length_fit_dict(traj, map_type='perlin_map', load=True)
    available_lengths = list(corr_len_fit_dict.values())
    closest_length = available_lengths[np.argmin(np.abs(np.array(available_lengths) - fit_corr_len))]
    closest_corresponding_scale = list(corr_len_fit_dict.keys())[list(corr_len_fit_dict.values()).index(closest_length)]
    if closest_length != fit_corr_len:
        print("Warning: desired fit correlation length {:.1f} not available. Taking {:.1f} instead".format(
            fit_corr_len, closest_length))
    return closest_corresponding_scale


def colorbar(mappable):
    from mpl_toolkits.axes_grid1 import make_axes_locatable
    import matplotlib.pyplot as plt
    last_axes = plt.gca()
    ax = mappable.axes
    fig = ax.figure
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    cbar = fig.colorbar(mappable, cax=cax)
    plt.sca(last_axes)
    return cbar


def plot_firing_rate_map_excitatory(traj, direction_idx, plot_run_names, exemplary_excitatory_neuron_id):
    max_val = 0
    for run_name in plot_run_names:
        fr_array = traj.results.runs[run_name].firing_rate_array
        f_rates = fr_array[:, direction_idx]
        run_max_val = np.max(f_rates)
        if run_max_val > max_val:
            # if traj.derived_parameters.runs[run_name].morphology.morph_label == POLARIZED:
            #     n_id_max_rate = np.argmax(f_rates)
            max_val = run_max_val

    n_id_polar_plot = exemplary_excitatory_neuron_id

    # Mark the neuron that is shown in polar plot
    ex_positions = traj.results.runs[plot_run_names[0]].ex_positions
    polar_plot_x, polar_plot_y = ex_positions[n_id_polar_plot]

    # Vertices for the plotted triangle
    tr_scale = 30.
    tr_x = tr_scale * np.cos(2. * np.pi / 3. + np.pi / 2.)
    tr_y = tr_scale * np.sin(2. * np.pi / 3. + np.pi / 2.) + polar_plot_y
    tr_vertices = np.array(
        [[polar_plot_x, polar_plot_y + tr_scale], [tr_x + polar_plot_x, tr_y], [-tr_x + polar_plot_x, tr_y]])

    height = 1 * panel_size
    width = 3 * panel_size

    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, (0.05, 0.15, 0.85, 0.7), axes_pad=panel_size / 6.0, cbar_location="right",
                     cbar_mode="single",
                     cbar_size="7%", cbar_pad=panel_size / 10.0,
                     nrows_ncols=(1, 3))

    for ax, run_name in zip(axes, [plot_run_names[i] for i in [2, 0, 1]]):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)
        firing_rate_array = traj.results[run_name].firing_rate_array
        number_of_excitatory_neurons_per_row = int(np.sqrt(traj.N_E))
        c = ax.pcolor(X, Y, np.reshape(firing_rate_array[:, direction_idx], (number_of_excitatory_neurons_per_row,
                                                                             number_of_excitatory_neurons_per_row)),
                      vmin=0, vmax=max_val, cmap='Reds')
        # ax.set_title(adjust_label(label))
        # ax.add_artist(Ellipse((polar_plot_x, polar_plot_y), 20., 20., color='k', fill=False, lw=2.))
        # ax.add_artist(Ellipse((polar_plot_x, polar_plot_y), 20., 20., color='w', fill=False, lw=1.))
        ax.add_artist(Polygon(tr_vertices, closed=True, fill=False, lw=2.5, color='k'))
        ax.add_artist(Polygon(tr_vertices, closed=True, fill=False, lw=1.5, color='w'))
        for spine in ax.spines.values():
            spine.set_edgecolor("grey")
            spine.set_linewidth(1)
        remove_ticks(ax)
    # fig.suptitle('spatial firing rate map', fontsize=16)
    ax.cax.colorbar(c)
    ax.cax.annotate("fr (Hz)", xy=(1, 1), xytext=(3, 3), xycoords="axes fraction", textcoords="offset points")
    # fig.tight_layout()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'C_firing_rate_map_excitatory' + FIGURE_SAVE_FORMAT, dpi=300)
        plt.close(fig)


def plot_firing_rate_map_inhibitory(traj, direction_idx, plot_run_names, selected_inhibitory_neuron):
    max_val = 0
    for run_name in plot_run_names:
        fr_array = traj.results.runs[run_name].inh_firing_rate_array
        f_rates = fr_array[:, direction_idx]
        run_max_val = np.max(f_rates)
        if run_max_val > max_val:
            max_val = run_max_val

    n_id_polar_plot = selected_inhibitory_neuron

    # Mark the neuron that is shown in Polar plot
    inhibitory_axonal_cloud_array = traj.results.runs[plot_run_names[1]].inhibitory_axonal_cloud_array
    polar_plot_x = inhibitory_axonal_cloud_array[n_id_polar_plot, 0]
    polar_plot_y = inhibitory_axonal_cloud_array[n_id_polar_plot, 1]

    width = 3 * panel_size
    height = panel_size

    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, (0.05, 0.15, 0.85, 0.7), axes_pad=panel_size / 6.0, cbar_location="right",
                     cbar_mode="single",
                     cbar_size="7%", cbar_pad=panel_size / 10.0,
                     nrows_ncols=(1, 3))

    for ax, run_name in zip(axes, [plot_run_names[i] for i in [2, 0, 1]]):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        inhibitory_axonal_cloud_array = traj.results.runs[run_name].inhibitory_axonal_cloud_array

        inh_positions = [[p[0], p[1]] for p in inhibitory_axonal_cloud_array]

        X, Y = get_position_mesh(inh_positions)
        inh_firing_rate_array = traj.results[run_name].inh_firing_rate_array
        number_of_inhibitory_neurons_per_row = int(np.sqrt(traj.N_I))
        c = ax.pcolor(X, Y, np.reshape(inh_firing_rate_array[:, direction_idx], (number_of_inhibitory_neurons_per_row,
                                                                                 number_of_inhibitory_neurons_per_row)),
                      vmin=0, vmax=max_val, cmap='Blues')
        # ax.set_title(adjust_label(label))

        circle_r = 40.
        ax.add_artist(Ellipse((polar_plot_x, polar_plot_y), circle_r, circle_r, color='k', fill=False, lw=4.5))
        ax.add_artist(Ellipse((polar_plot_x, polar_plot_y), circle_r, circle_r, color='w', fill=False, lw=3))
        for spine in ax.spines.values():
            spine.set_edgecolor("grey")
            spine.set_linewidth(0.5)

        remove_ticks(ax)
        # fig.colorbar(c, ax=ax, label="f (Hz)")
    # fig.suptitle('spatial firing rate map', fontsize=16)
    cbar = ax.cax.colorbar(c)
    cbar_ticks = range(0,151,30)
    cbar.ax.set_yticks(cbar_ticks)
    ax.cax.annotate("fr (Hz)", xy=(1, 1), xytext=(3, 3), xycoords="axes fraction", textcoords="offset points")
    # fig.tight_layout()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'C_firing_rate_map_inhibitory' + FIGURE_SAVE_FORMAT, dpi=300)
        plt.close(fig)
    return max_val


def normal_labels(label):
    if label == POLARIZED:
        label = 'polarized'
    elif label == NO_SYNAPSES:
        label = 'no interneurons'
    return label


def short_labels(label):
    if label == POLARIZED:
        label = 'pol.'
    elif label == CIRCULAR:
        label = 'cir.'
    elif label == "no conn":
        label = "no inh."
    return label


def plot_input_map(traj, run_name, figsize=(panel_size, panel_size), figname='input_map' + FIGURE_SAVE_FORMAT):
    n_ex = int(np.sqrt(traj.N_E))

    width, height = figsize

    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, (0.15, 0.1, 0.7, 0.8), axes_pad=panel_size / 3.0, cbar_location="right", cbar_mode=None,
                     nrows_ncols=(1, 1))
    ax = axes[0]

    traj.f_set_crun(run_name)

    X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)
    head_dir_preference = np.array(traj.results.runs[run_name].ex_tunings).reshape((n_ex, n_ex))
    scale_length = traj.morphology.long_axis * 2
    traj.f_restore_default()

    ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap=head_direction_input_colormap)

    for spine in ax.spines.values():
        spine.set_edgecolor("grey")
        spine.set_linewidth(0.5)
    remove_ticks(ax)

    start_scale_x = 100
    end_scale_x = start_scale_x + scale_length
    start_scale_y = -70
    end_scale_y = start_scale_y

    add_length_scale(ax, scale_length, start_scale_x, end_scale_x, start_scale_y, end_scale_y)
    ax.cax.set_visible(False)

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + figname)
        plt.close(fig)


def plot_example_input_maps(traj, figsize=(panel_size, panel_size), figname='perlin_example_input_maps' + FIGURE_SAVE_FORMAT):
    n_ex = int(np.sqrt(traj.N_E))

    width, height = figsize

    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, (0.1, 0.1, 0.85, 0.85), axes_pad=panel_size / 12.0, nrows_ncols=(3, 3))

    fit_corr_len_list = [20, 50, 100]
    seed_list = [1, 2, 3]
    corr_and_seed = itertools.product(fit_corr_len_list, seed_list)

    corr_len_fit_dict = correlation_length_fit_dict(traj, map_type='perlin_map', load=True)

    for idx, (ax, (fit_corr_len, seed)) in enumerate(zip(axes, corr_and_seed)):
        corresp_scale = get_closest_fit_correlation_length(traj, fit_corr_len)
        par_dict = {'seed': seed, 'scale': corresp_scale}

        run_name = filter_run_names_by_par_dict(traj, par_dict)[0]
        traj.f_set_crun(run_name)

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)
        head_dir_preference = np.array(traj.results.runs[run_name].ex_tunings).reshape((n_ex, n_ex))
        scale_length = traj.morphology.long_axis * 2
        traj.f_restore_default()

        ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap=head_direction_input_colormap)

        for spine in ax.spines.values():
            spine.set_edgecolor("grey")
            spine.set_linewidth(0.5)
        remove_ticks(ax)

        ax.set_ylabel('{:3.0f} um'.format(corr_len_fit_dict[corresp_scale]))

        inhibitory_axonal_cloud_array = traj.results.runs[run_name].inhibitory_axonal_cloud_array

        plt_polar_id = 255
        plt_polar_coordinates = inhibitory_axonal_cloud_array[plt_polar_id]
        plt_polar_neuron = Interneuron(plt_polar_coordinates[0],
                                  plt_polar_coordinates[1],
                                  traj.morphology.long_axis,
                                  traj.morphology.short_axis,
                                  plt_polar_coordinates[2])

        circ_radius = np.sqrt(traj.morphology.long_axis * traj.morphology.short_axis)

        plt_circ_id = 65
        plt_circ_coordinates = inhibitory_axonal_cloud_array[plt_circ_id]
        plt_circ_neuron = Interneuron(plt_circ_coordinates[0],
                                 plt_circ_coordinates[1],
                                 circ_radius,
                                 circ_radius,
                                 plt_circ_coordinates[2])

        plt_interneurons = [plt_polar_neuron, plt_circ_neuron]

        for p in plt_interneurons:
            ell = p.get_ellipse()
            edgecolor = 'black'
            alpha = 1
            zorder = 10
            linewidth = 1.

            ell.set_edgecolor(edgecolor)
            ell.set_alpha(alpha)
            ell.set_zorder(zorder)
            ell.set_linewidth(linewidth)
            ax.add_artist(ell)

        if idx == 8:
            start_scale_x = 100
            end_scale_x = start_scale_x + scale_length
            start_scale_y = -70
            end_scale_y = start_scale_y

            add_length_scale(ax, scale_length, start_scale_x, end_scale_x, start_scale_y, end_scale_y)
            ax.cax.set_visible(False)

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + figname)
        plt.close(fig)


def plot_axonal_clouds(traj, plot_run_names):
    n_ex = int(np.sqrt(traj.N_E))

    height = 1 * panel_size
    width = 3 * panel_size

    cluster_positions = [(250, 250), (750, 200), (450, 600)]
    cluster_sizes = [4, 4, 4]
    selected_neurons = []

    inhibitory_positions = traj.results.runs[plot_run_names[0]].inhibitory_axonal_cloud_array[:, :2]
    for cluster_position, number_of_neurons_in_cluster in zip(cluster_positions, cluster_sizes):
        selection = get_neurons_close_to_given_position(cluster_position, number_of_neurons_in_cluster,
                                                        inhibitory_positions)
        selected_neurons.extend(selection)

    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, 111, axes_pad=0.15, cbar_location="right", cbar_mode="single", cbar_size="7%",
                     nrows_ncols=(1, 3))
    for ax, run_name in zip(axes, [plot_run_names[i] for i in [2, 0, 1]]):
        traj.f_set_crun(run_name)

        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)

        inhibitory_axonal_cloud_array = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        axonal_clouds = [Interneuron(p[0], p[1], traj.morphology.long_axis, traj.morphology.short_axis, p[2]) for p in
                         inhibitory_axonal_cloud_array]

        head_dir_preference = np.array(traj.results.runs[run_name].ex_tunings).reshape((n_ex, n_ex))
        c = ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap=head_direction_input_colormap)
        # ax.set_title(normal_labels(label))
        ax.set_aspect('equal')
        remove_frame(ax)
        remove_ticks(ax)
        # fig.colorbar(c, ax=ax, label="Tuning")

        if label != NO_SYNAPSES and axonal_clouds is not None:
            for i, p in enumerate(axonal_clouds):
                ell = p.get_ellipse()
                if i in selected_neurons:
                    edgecolor = 'black'
                    alpha = 1
                    zorder = 10
                    linewidth = 1
                else:
                    edgecolor = 'gray'
                    alpha = 0.5
                    zorder = 1
                    linewidth = 0.3

                ell.set_edgecolor(edgecolor)
                ell.set_alpha(alpha)
                ell.set_zorder(zorder)
                ell.set_linewidth(linewidth)
                ax.add_artist(ell)

    traj.f_set_crun(plot_run_names[0])
    scale_length = traj.morphology.long_axis * 2
    traj.f_restore_default()

    start_scale_x = 100
    end_scale_x = start_scale_x + scale_length
    start_scale_y = -70
    end_scale_y = start_scale_y

    add_length_scale(axes[0], scale_length, start_scale_x, end_scale_x, start_scale_y, end_scale_y)

    axes[1].set_yticklabels([])
    axes[2].set_yticklabels([])
    axes[0].cax.set_visible(False)
    traj.f_restore_default()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'B_i_axonal_clouds' + FIGURE_SAVE_FORMAT)
        plt.close(fig)


def get_neurons_close_to_given_position(cluster_position, number_of_neurons_in_cluster, positions):
    position = np.array(cluster_position)
    distance_vectors = positions - np.expand_dims(position, 0).repeat(positions.shape[0],
                                                                      axis=0)
    distances = np.linalg.norm(distance_vectors, axis=1)
    selection = list(np.argpartition(distances, number_of_neurons_in_cluster)[:number_of_neurons_in_cluster])
    return selection


def plot_orientation_maps_diff_scales(traj):
    n_ex = int(np.sqrt(traj.N_E))

    scale_run_names = []
    plot_scales = [0.0, 100.0, 200.0, 300.0]
    for scale in plot_scales:
        par_dict = {'seed': 1, 'correlation_length': get_closest_scale(traj, scale), 'long_axis': 100.}
        scale_run_names.append(*filter_run_names_by_par_dict(traj, par_dict))

    fig, axes = plt.subplots(1, 4, figsize=(18., 4.5))
    for ax, run_name, scale in zip(axes, scale_run_names, plot_scales):
        traj.f_set_crun(run_name)

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)

        head_dir_preference = np.array(traj.results.runs[run_name].ex_tunings).reshape((n_ex, n_ex))
        # TODO: Why was this transposed for plotting? (now changed)
        c = ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap='twilight')
        ax.set_title('Correlation length: {}'.format(scale))
        fig.colorbar(c, ax=ax, label="Tuning")

    # fig.suptitle('axonal cloud', fontsize=16)
    traj.f_restore_default()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'orientation_maps_diff_scales' + FIGURE_SAVE_FORMAT)
        plt.close(fig)


def plot_polar_plot_excitatory(traj, plot_run_names, selected_neuron_idx,
                               figname=FIGURE_SAVE_PATH + 'D_polar_plot_excitatory' + FIGURE_SAVE_FORMAT):
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    directions_plt = list(directions)
    directions_plt.append(directions[0])
    height = panel_size
    width = panel_size
    fig, ax = plt.subplots(1, 1, figsize=(height, width), subplot_kw=dict(projection='polar'))
    # head_direction_indices = traj.results.runs[plot_run_names[0]].head_direction_indices
    # sorted_ids = np.argsort(head_direction_indices)
    # plot_n_idx = sorted_ids[-75]
    plot_n_idx = selected_neuron_idx

    line_styles = ['dotted', 'solid', 'dashed']
    colors = ['r', 'lightsalmon', 'grey']
    labels = ['pol.   ', 'cir.               ', 'no inh.']
    line_widths = [1.5, 1.5, 1]
    zorders = [10, 2, 1]

    max_rate = 0.0
    ax.plot([], [], color='white',label='          ')
    hdis = []
    for run_idx, run_name in enumerate(plot_run_names):
        label = labels[run_idx]
        hdi = traj.results.runs[run_name].head_direction_indices[selected_neuron_idx]
        hdis.append(hdi)
        tuning_vectors = traj.results.runs[run_name].tuning_vectors
        rate_plot = [np.linalg.norm(v) for v in tuning_vectors[plot_n_idx]]
        run_max_rate = np.max(rate_plot)
        if run_max_rate > max_rate:
            max_rate = run_max_rate
        rate_plot.append(rate_plot[0])
        plt_label = '{:s}'.format(short_labels(label))
        ax.plot(directions_plt, rate_plot, linewidth=line_widths[run_idx],
                label=plt_label, color=colors[run_idx], linestyle=line_styles[run_idx], zorder=zorders[run_idx])
    ticks = [40., 80.]
    ax.set_rgrids(ticks, labels=["{:.0f} Hz".format(ticklabel) if idx == len(ticks) - 1 else "" for idx, ticklabel in
                                 enumerate(ticks)], angle=60)
    ax.set_thetagrids([0, 90, 180, 270], labels=[])
    ax.xaxis.grid(linewidth=0.4)
    ax.yaxis.grid(linewidth=0.4)
    leg = ax.legend(loc="lower right", bbox_to_anchor=(1.15, -0.2), handlelength=1, fontsize="medium")
    leg.get_frame().set_linewidth(0.0)

    hdi_box_x, hdi_box_y = (0.86, -0.09)
    hdi_box_dy = 0.14
    hdi_box = ax.text(hdi_box_x, hdi_box_y + 3 * hdi_box_dy, 'HDI', fontsize='medium', transform=ax.transAxes,zorder=9.)
    hdi_box = ax.text(hdi_box_x, hdi_box_y + 2 * hdi_box_dy, '{:.2f}'.format(hdis[0]), fontsize='medium',transform=ax.transAxes, zorder=9.)
    hdi_box = ax.text(hdi_box_x, hdi_box_y + hdi_box_dy, '{:.2f}'.format(hdis[1]), fontsize='medium', transform=ax.transAxes,zorder=9.)
    hdi_box = ax.text(hdi_box_x, hdi_box_y, '{:.2f}'.format(hdis[2]), fontsize='medium', transform=ax.transAxes,zorder=9.)

    ax.axes.spines["polar"].set_visible(False)
    if save_figs:
        plt.savefig(figname)
        plt.close(fig)


def plot_polar_plot_inhibitory(traj, plot_run_names, selected_neuron_idx, figname=FIGURE_SAVE_PATH +
                                                                                  'D_polar_plot_inhibitory' + FIGURE_SAVE_FORMAT):
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    directions_plt = list(directions)
    directions_plt.append(directions[0])
    height = panel_size
    width = panel_size
    fig, ax = plt.subplots(1, 1, figsize=(height, width), subplot_kw=dict(projection='polar'))
    # head_direction_indices = traj.results.runs[plot_run_names[0]].inh_head_direction_indices
    # sorted_ids = np.argsort(head_direction_indices)
    # plot_n_idx = sorted_ids[-75]
    plot_n_idx = selected_neuron_idx

    line_styles = ['dotted', 'solid']
    colors = ['b', 'lightblue']
    labels = ['pol. ', 'cir.            ']

    line_widths = [1.5, 1.5]
    zorders = [10, 2]

    ax.plot([], [], color='white',label='          ')
    hdis = []
    for run_idx, run_name in enumerate(plot_run_names[:2]):
        label = labels[run_idx]
        hdi = traj.results.runs[run_name].inh_head_direction_indices[selected_neuron_idx]
        hdis.append(hdi)

        tuning_vectors = traj.results.runs[run_name].inh_tuning_vectors
        rate_plot = [np.linalg.norm(v) for v in tuning_vectors[plot_n_idx]]
        rate_plot.append(rate_plot[0])
        plt_label = '{:s}'.format(short_labels(label))
        ax.plot(directions_plt, rate_plot, linewidth=line_widths[run_idx],
                label=plt_label,
                color=colors[run_idx], linestyle=line_styles[run_idx], zorder=zorders[run_idx])
    ticks = [40., 80., 120.]
    ax.set_rgrids(ticks, labels=["{:.0f} Hz".format(ticklabel) if idx == len(ticks) - 1 else "" for idx, ticklabel in
                                 enumerate(ticks)], angle=60)
    ax.set_thetagrids([0, 90, 180, 270], labels=[])
    ax.xaxis.grid(linewidth=0.4)
    ax.yaxis.grid(linewidth=0.4)
    leg = ax.legend(loc="lower right", bbox_to_anchor=(1.15, -0.15), handlelength=1, fontsize="medium")
    leg.get_frame().set_linewidth(0.0)

    hdi_box_x, hdi_box_y = (0.86, -0.04)
    hdi_box_dy = 0.14
    hdi_box = ax.text(hdi_box_x, hdi_box_y+2*hdi_box_dy, 'HDI', fontsize='medium', transform=ax.transAxes, zorder=9.)
    hdi_box = ax.text(hdi_box_x, hdi_box_y+hdi_box_dy, '{:.2f}'.format(hdis[0]), fontsize='medium', transform=ax.transAxes, zorder=9.)
    hdi_box = ax.text(hdi_box_x, hdi_box_y, '{:.2f}'.format(hdis[1]), fontsize='medium', transform=ax.transAxes, zorder=9.)


    ax.axes.spines["polar"].set_visible(False)

    if save_figs:
        plt.savefig(figname)
        plt.close(fig)


def plot_hdi_histogram_combined_and_overlayed(traj, plot_run_names, ex_polar_plot_id, in_polar_plot_id, cut_off_dist):
    labels = []
    inh_hdis = []
    exc_hdis = []
    no_conn_hdi = 0.

    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        if label != NO_SYNAPSES:
            labels.append(normal_labels(label))

            inh_head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
            inh_axonal_cloud = traj.results.runs[run_name].inhibitory_axonal_cloud_array
            inh_cut_off_ids = (inh_axonal_cloud[:, 0] >= cut_off_dist) & \
                              (inh_axonal_cloud[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                              (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                              (inh_axonal_cloud[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
            inh_hdis.append(sorted(inh_head_direction_indices[inh_cut_off_ids]))

            exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
            ex_positions = traj.results.runs[run_name].ex_positions
            exc_cut_off_ids = (ex_positions[:, 0] >= cut_off_dist) & \
                              (ex_positions[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                              (ex_positions[:, 1] >= cut_off_dist) & \
                              (ex_positions[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
            exc_hdis.append(sorted(exc_head_direction_indices[exc_cut_off_ids]))

        else:
            exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
            no_conn_hdi = np.mean(exc_head_direction_indices)

    # Look for a representative excitatory neuron
    hdi_mean_dict = {}
    excitatory_hdi_means = [np.mean(hdis) for hdis in exc_hdis]
    hdi_mean_dict["polar_exc"] = excitatory_hdi_means[0]
    hdi_mean_dict["circular_exc"] = excitatory_hdi_means[1]

    inhibitory_hdi_means = [np.mean(hdis) for hdis in inh_hdis]
    hdi_mean_dict["polar_inh"] = inhibitory_hdi_means[0]
    hdi_mean_dict["circular_inh"] = inhibitory_hdi_means[1]

    width = 2 * panel_size
    height = 1.2 * panel_size
    fig, axes = plt.subplots(2, 1, figsize=(width, height))
    plt.subplots_adjust(wspace=0, hspace=0.1)
    bins = np.linspace(0.0, 1.0, 21, endpoint=True)

    max_density = 0
    for i in range(2):
        # i = i + 1 # grid spec indexes from 0
        # ax = fig.add_subplot(gs[i])
        ax = axes[i]

        density_e, _, _ = ax.hist(exc_hdis[i], color='r', edgecolor='r', alpha=0.3, bins=bins, density=True)
        density_i, _, _ = ax.hist(inh_hdis[i], color='b', edgecolor='b', alpha=0.3, bins=bins, density=True)

        max_density = np.max([max_density, np.max(density_e), np.max(density_i)])

        ax.axvline(np.mean(exc_hdis[i]), color='r')

        ax.axvline(np.mean(inh_hdis[i]), color='b')
        ax.axvline(no_conn_hdi, color='dimgrey', linestyle='--', linewidth=1.5)

        ax.set_ylabel(labels[i], rotation='vertical')
        # plt.axis('on')
        if i == 0:
            ax.set_xticklabels([])

        else:
            ax.set_xlabel('HDI')
        remove_frame(ax, ["top", "right", "bottom"])
    max_density = 1.2 * max_density
    fig.subplots_adjust(left=0.2, right=0.95, bottom=0.2)

    axes[0].annotate('% cells', (0, 1.0), xycoords='axes fraction', va="bottom", ha="right")

    axes[1].annotate("no ihn.\n{:.2f}".format(no_conn_hdi), xy=(no_conn_hdi, max_density),
                     xytext=(-2, 0), xycoords="data",
                     textcoords="offset points",
                     va="top", ha="right", color="dimgrey")

    for i, ax in enumerate(axes):
        ax.annotate("{:.2f}".format(np.mean(exc_hdis[i])), xy=(np.mean(exc_hdis[i]), max_density),
                    xytext=(2, 0), xycoords="data",
                    textcoords="offset points",
                    va="top", ha="left", color="r")

        # i_ha = "left" if i == 1 else "right"
        # i_offset = 2 if i == 1 else -2
        i_ha = "right"
        i_offset = -1
        ax.annotate("{:.2f}".format(np.mean(inh_hdis[i])), xy=(np.mean(inh_hdis[i]), max_density),
                    xytext=(i_offset, 0), xycoords="data",
                    textcoords="offset points",
                    va="top", ha=i_ha, color="b")

    for ax in axes:
        ax.set_ylim(0, max_density)
    # plt.annotate('probability density', (-0.2,1.5), xycoords='axes fraction', rotation=90, fontsize=18)

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'E_hdi_histogram_combined_and_overlayed' + FIGURE_SAVE_FORMAT.format(cut_off_dist))
        plt.close(fig)

    return hdi_mean_dict


def get_neurons_with_given_hdi(polar_hdi, circular_hdi, max_number_of_suggestions, plot_run_names, traj, type):
    polar_run_name = plot_run_names[0]
    circular_run_name = plot_run_names[1]
    polar_ex_hdis = traj.results.runs[polar_run_name].head_direction_indices if type == "ex" else traj.results.runs[
        polar_run_name].inh_head_direction_indices
    circular_ex_hdis = traj.results.runs[circular_run_name].head_direction_indices if type == "ex" else \
        traj.results.runs[
            polar_run_name].inh_head_direction_indices
    neuron_indices = get_indices_of_closest_values(polar_ex_hdis, polar_hdi,
                                                   circular_ex_hdis,
                                                   circular_hdi, 0.1 * np.abs(
            polar_hdi - circular_hdi), max_number_of_suggestions)
    return neuron_indices


def get_indices_of_closest_values(first_list, first_value, second_list, second_value, absolute_tolerance_list_one,
                                  number_of_indices):
    is_close_in_list_one = np.abs(first_list - first_value) < absolute_tolerance_list_one
    indices_close_in_list_one = np.where(is_close_in_list_one)[0]
    indices_closest_in_list_two = indices_close_in_list_one[np.argpartition(np.abs(second_list[
                                                                                       indices_close_in_list_one] - second_value),
                                                                            number_of_indices)]
    return indices_closest_in_list_two[:number_of_indices]


def filter_run_names_by_par_dict(traj, par_dict):
    run_name_list = []
    for run_idx, run_name in enumerate(traj.f_get_run_names()):
        traj.f_set_crun(run_name)
        paramters_equal = True
        for key, val in par_dict.items():
            if (traj.par[key] != val):
                paramters_equal = False
        if paramters_equal:
            run_name_list.append(run_name)
    traj.f_restore_default()

    return run_name_list


def plot_exc_and_inh_hdi_over_simplex_grid_scale(traj, plot_run_names, cut_off_dist):
    corr_len_expl = traj.f_get('scale').f_get_range()
    seed_expl = traj.f_get('seed').f_get_range()
    label_expl = [traj.derived_parameters.runs[run_name].morphology.morph_label for run_name in traj.f_get_run_names()]
    label_range = set(label_expl)

    exc_hdi_frame = pd.Series(index=[corr_len_expl, seed_expl, label_expl])
    exc_hdi_frame.index.names = ["corr_len", "seed", "label"]
    inh_hdi_frame = pd.Series(index=[corr_len_expl, seed_expl, label_expl])
    inh_hdi_frame.index.names = ["corr_len", "seed", "label"]

    for run_name, corr_len, seed, label in zip(plot_run_names, corr_len_expl, seed_expl, label_expl):
        ex_tunings = traj.results.runs[run_name].ex_tunings
        inh_hdis = []
        exc_hdis = []

        inh_head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
        inh_axonal_cloud = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        inh_cut_off_ids = (inh_axonal_cloud[:, 0] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
        inh_hdis.append(sorted(inh_head_direction_indices[inh_cut_off_ids]))

        exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
        ex_positions = traj.results.runs[run_name].ex_positions
        exc_cut_off_ids = (ex_positions[:, 0] >= cut_off_dist) & \
                          (ex_positions[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                          (ex_positions[:, 1] >= cut_off_dist) & \
                          (ex_positions[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
        exc_hdis.append(sorted(exc_head_direction_indices[exc_cut_off_ids]))

        exc_hdi_frame[corr_len, seed, label] = np.mean(exc_hdis)

        inh_hdi_frame[corr_len, seed, label] = np.mean(inh_hdis)

    # TODO: Standard deviation also for the population
    exc_hdi_n_and_seed_mean = exc_hdi_frame.groupby(level=[0, 2]).mean()
    exc_hdi_n_and_seed_std_dev = exc_hdi_frame.groupby(level=[0, 2]).std()

    inh_hdi_n_and_seed_mean = inh_hdi_frame.groupby(level=[0, 2]).mean()
    inh_hdi_n_and_seed_std_dev = inh_hdi_frame.groupby(level=[0, 2]).std()

    markersize = 4.
    exc_style_dict = {
        NO_SYNAPSES: ['dimgrey', 'dashed', '', 0],
        POLARIZED: ['red', 'solid', '^', markersize],
        CIRCULAR: ['lightsalmon', 'solid', '^', markersize]
    }
    inh_style_dict = {
        NO_SYNAPSES: ['dimgrey', 'dashed', '', 0],
        POLARIZED: ['blue', 'solid', 'o', markersize],
        CIRCULAR: ['lightblue', 'solid', 'o', markersize]

    }

    width = 2 * panel_size
    height = 1.2 * panel_size
    fig, ax = plt.subplots(1, 1, figsize=(width, height))
    for label in sorted(label_range, reverse=True):
        if label == NO_SYNAPSES:
            no_conn_hdi = exc_hdi_n_and_seed_mean[get_closest_scale(traj, 200), label]
            ax.axhline(no_conn_hdi, color='grey', linestyle='--')

            ax.annotate(short_labels(label), xy=(1.0, no_conn_hdi), xytext=(0, -2), xycoords='axes fraction',
                        textcoords="offset points",
                        va="top", \
                        ha="right",
                        color="dimgrey")

            continue
        exc_hdi_mean = exc_hdi_n_and_seed_mean[:, label]
        exc_hdi_std = exc_hdi_n_and_seed_std_dev[:, label]

        inh_hdi_mean = inh_hdi_n_and_seed_mean[:, label]
        inh_hdi_std = inh_hdi_n_and_seed_std_dev[:, label]

        corr_len_range = exc_hdi_mean.keys().to_numpy()

        exc_col, exc_lin, exc_mar, exc_mar_size = exc_style_dict[label]
        inh_col, inh_lin, inh_mar, inh_mar_size = inh_style_dict[label]

        simplex_grid_scale = corr_len_range * np.sqrt(2)

        ax.plot(simplex_grid_scale, exc_hdi_mean, label='exc., ' + label, marker=exc_mar, color=exc_col, linestyle=exc_lin,
                markersize=exc_mar_size, alpha=0.5)
        plt.fill_between(simplex_grid_scale, exc_hdi_mean - exc_hdi_std,
                         exc_hdi_mean + exc_hdi_std, alpha=0.3, color=exc_col)
        ax.plot(simplex_grid_scale, inh_hdi_mean, label='inh., ' + label, marker=inh_mar, color=inh_col, linestyle=inh_lin,
                markersize=inh_mar_size, alpha=0.5)
        plt.fill_between(simplex_grid_scale, inh_hdi_mean - inh_hdi_std,
                         inh_hdi_mean + inh_hdi_std, alpha=0.3, color=inh_col)

    ax.set_xlabel('simplex grid scale')
    ax.set_ylabel('HDI')
    ax.axvline(get_closest_scale(traj, 200.0) * np.sqrt(2), color='k', linewidth=0.5, zorder=0)
    ax.set_ylim(0.0, 1.0)
    # ax.set_xlim(0.0, np.max(corr_len_range))
    remove_frame(ax, ["right", "top"])

    tablelegend(ax, ncol=2, bbox_to_anchor=(1.1, 1.1), loc="upper right",
                row_labels=None,
                col_labels=[short_labels(label) for label in sorted(label_range - {"no conn"}, reverse=True)],
                title_label='', borderaxespad=0, handlelength=2, edgecolor='white')
    fig.subplots_adjust(bottom=0.2, left=0.2)
    # plt.legend()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'F_hdi_over_grid_scale' + FIGURE_SAVE_FORMAT)
        plt.close(fig)


def plot_exc_and_inh_hdi_over_fit_corr_len(traj, plot_run_names, cut_off_dist):
    corr_len_expl = traj.f_get('scale').f_get_range()
    seed_expl = traj.f_get('seed').f_get_range()
    label_expl = [traj.derived_parameters.runs[run_name].morphology.morph_label for run_name in traj.f_get_run_names()]
    label_range = set(label_expl)

    exc_hdi_frame = pd.Series(index=[corr_len_expl, seed_expl, label_expl])
    exc_hdi_frame.index.names = ["corr_len", "seed", "label"]
    inh_hdi_frame = pd.Series(index=[corr_len_expl, seed_expl, label_expl])
    inh_hdi_frame.index.names = ["corr_len", "seed", "label"]

    for run_name, corr_len, seed, label in zip(plot_run_names, corr_len_expl, seed_expl, label_expl):
        ex_tunings = traj.results.runs[run_name].ex_tunings

        inh_hdis = []
        exc_hdis = []

        inh_head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
        inh_axonal_cloud = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        inh_cut_off_ids = (inh_axonal_cloud[:, 0] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
        inh_hdis.append(sorted(inh_head_direction_indices[inh_cut_off_ids]))

        exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
        ex_positions = traj.results.runs[run_name].ex_positions
        exc_cut_off_ids = (ex_positions[:, 0] >= cut_off_dist) & \
                          (ex_positions[:, 0] <= traj.parameters.input_map.sheet_size - cut_off_dist) & \
                          (ex_positions[:, 1] >= cut_off_dist) & \
                          (ex_positions[:, 1] <= traj.parameters.input_map.sheet_size - cut_off_dist)
        exc_hdis.append(sorted(exc_head_direction_indices[exc_cut_off_ids]))

        exc_hdi_frame[corr_len, seed, label] = np.mean(exc_hdis)

        inh_hdi_frame[corr_len, seed, label] = np.mean(inh_hdis)

    # TODO: Standard deviation also for the population
    exc_hdi_n_and_seed_mean = exc_hdi_frame.groupby(level=[0, 2]).mean()
    exc_hdi_n_and_seed_std_dev = exc_hdi_frame.groupby(level=[0, 2]).std()

    inh_hdi_n_and_seed_mean = inh_hdi_frame.groupby(level=[0, 2]).mean()
    inh_hdi_n_and_seed_std_dev = inh_hdi_frame.groupby(level=[0, 2]).std()

    markersize = 4.
    exc_style_dict = {
        NO_SYNAPSES: ['dimgrey', 'dashed', '', 0],
        POLARIZED: ['red', 'solid', '^', markersize],
        CIRCULAR: ['lightsalmon', 'solid', '^', markersize]
    }
    inh_style_dict = {
        NO_SYNAPSES: ['dimgrey', 'dashed', '', 0],
        POLARIZED: ['blue', 'solid', 'o', markersize],
        CIRCULAR: ['lightblue', 'solid', 'o', markersize]

    }
    # colors = ['blue', 'grey', 'lightblue']
    # linestyles = ['solid', 'dashed', 'solid']
    # markers = [verts, '', 'o']

    corr_len_fit_dict = correlation_length_fit_dict(traj, map_type='perlin_map', load=True)

    width = 2 * panel_size
    height = 1.2 * panel_size
    fig, ax = plt.subplots(1, 1, figsize=(width, height))

    for label in sorted(label_range, reverse=True):
        if label == NO_SYNAPSES:
            no_conn_hdi = exc_hdi_n_and_seed_mean[get_closest_scale(traj, 200), label]
            ax.axhline(no_conn_hdi, color='grey', linestyle='--')

            ax.annotate(short_labels(label), xy=(1.0, no_conn_hdi), xytext=(0, -2), xycoords='axes fraction',
                        textcoords="offset points",
                        va="top", \
                        ha="right",
                        color="dimgrey")

            continue
        exc_hdi_mean = exc_hdi_n_and_seed_mean[:, label]
        exc_hdi_std = exc_hdi_n_and_seed_std_dev[:, label]

        inh_hdi_mean = inh_hdi_n_and_seed_mean[:, label]
        inh_hdi_std = inh_hdi_n_and_seed_std_dev[:, label]

        corr_len_range = exc_hdi_mean.keys().to_numpy()

        exc_col, exc_lin, exc_mar, exc_mar_size = exc_style_dict[label]
        inh_col, inh_lin, inh_mar, inh_mar_size = inh_style_dict[label]

        fit_corr_len = [corr_len_fit_dict[corr_len] for corr_len in corr_len_range]

        ax.plot(fit_corr_len, exc_hdi_mean, label='exc., ' + label, marker=exc_mar, color=exc_col, linestyle=exc_lin,
                markersize=exc_mar_size, alpha=0.5)
        plt.fill_between(fit_corr_len, exc_hdi_mean - exc_hdi_std,
                         exc_hdi_mean + exc_hdi_std, alpha=0.3, color=exc_col)
        ax.plot(fit_corr_len, inh_hdi_mean, label='inh., ' + label, marker=inh_mar, color=inh_col, linestyle=inh_lin,
                markersize=inh_mar_size, alpha=0.5)
        plt.fill_between(fit_corr_len, inh_hdi_mean - inh_hdi_std,
                         inh_hdi_mean + inh_hdi_std, alpha=0.3, color=inh_col)

    ax.set_xlabel('correlation length (um)')
    ax.set_ylabel('HDI')
    ax.axvline(corr_len_fit_dict[get_closest_scale(traj, 200.0)], color='k', linewidth=0.5, zorder=0)
    ax.set_ylim(0.0, 1.0)
    remove_frame(ax, ["right", "top"])

    tablelegend(ax, ncol=2, bbox_to_anchor=(1.1, 1.1), loc="upper right",
                row_labels=None,
                col_labels=[short_labels(label) for label in sorted(label_range - {"no conn"}, reverse=True)],
                title_label='', borderaxespad=0, handlelength=2, edgecolor='white')
    fig.subplots_adjust(bottom=0.2, left=0.2)
    # plt.legend()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'F_hdi_over_corr_len_scaled' + FIGURE_SAVE_FORMAT)
        plt.close(fig)


def get_phase_difference(total_difference):
    """
    Map accumulated phase difference to shortest possible difference
    :param total_difference:
    :return: relative_difference
    """
    return (total_difference + np.pi) % (2 * np.pi) - np.pi


def plot_firing_rate_similar_vs_diff_tuning(traj, plot_run_names, figsize=(9, 9)):
    # The plot that Imre wanted
    n_bins = traj.parameters.input.number_of_directions
    fig, ax = plt.subplots(1, 1, figsize=figsize)
    dir_bins = np.linspace(-np.pi, np.pi, n_bins + 1)

    plot_fr_array = []
    labels = []

    similarity_threshold = np.pi / 6.
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    for run_idx, run_name in enumerate(plot_run_names):

        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        labels.append(short_labels(label))
        fr_similar_tunings = []
        fr_different_tunings = []

        ex_tunings = traj.results.runs[run_name].ex_tunings
        firing_rate_array = traj.results[run_name].firing_rate_array

        for tuning, firing_rates in zip(ex_tunings, firing_rate_array):
            for idx, dir in enumerate(directions):
                if np.abs(get_phase_difference(tuning - dir)) <= similarity_threshold:
                    fr_similar_tunings.append(firing_rates[idx])
                elif np.abs(get_phase_difference(tuning + np.pi - dir)) <= similarity_threshold:
                    fr_different_tunings.append(firing_rates[idx])
        plot_fr_array.append([np.mean(fr_similar_tunings), np.mean(fr_different_tunings)])

    x = np.arange(3)  # the label locations
    width = 0.35  # the width of the bars

    plot_fr_array = np.array(plot_fr_array)

    rects1 = ax.bar(x - width / 2, plot_fr_array[:, 0], width,
                    label=r'preferred')
    rects2 = ax.bar(x + width / 2, plot_fr_array[:, 1], width,
                    label=r'opposite')

    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    # ax.set_title('Mean firing rate for tunings similar and different to input')
    # ax.set_ylabel('Mean firing rate')
    ax.annotate(r'$\overline{\mathrm{fr}}$ (Hz)', (0.05, 1.0), xycoords='axes fraction', va="bottom", ha="right")
    leg = ax.legend(loc="upper left", bbox_to_anchor=(0.0, 1.2), handlelength=1, fontsize="medium")
    leg.get_frame().set_linewidth(0.0)

    def autolabel(rects):
        """Attach a text label above each bar in *rects*, displaying its height."""
        for rect in rects:
            height = rect.get_height()
            ax.annotate('{}'.format(np.round(height)),
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom')

    autolabel(rects1)
    autolabel(rects2)

    fig.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'SUPPLEMENT_B_firing_rate_similar_vs_diff_tuning' + FIGURE_SAVE_FORMAT, dpi=200)
        plt.close(fig)


def get_firing_rates_along_preferred_axis(traj, run_name, neuron_idx):
    firing_rates = traj.results[run_name].firing_rate_array[neuron_idx, :]
    tuning = traj.results[run_name].ex_tunings[neuron_idx]
    anti_tuning = tuning + np.pi if tuning + np.pi < np.pi else tuning - np.pi

    tuning_idx = np.argmin(np.abs(directions - tuning))
    anti_tuning_idx = np.argmin(np.abs(directions - anti_tuning))

    firing_at_the_preferred_direction = firing_rates[tuning_idx]
    firing_at_the_opposite_direction = firing_rates[anti_tuning_idx]

    return firing_at_the_preferred_direction, firing_at_the_opposite_direction


def get_hdi(traj, run_name, neuron_idx, type):
    return traj.results.runs[run_name].head_direction_indices[neuron_idx] if type=="ex" else traj.results.runs[
        run_name].inh_head_direction_indices[neuron_idx]


def plot_colorbar(figsize=(2, 2), figname=None):
    azimuth_no = 360
    zenith_no = 15

    azimuths = np.linspace(-180, 180, azimuth_no)
    zeniths = np.linspace(0.85, 1, zenith_no)

    values = azimuths * np.ones((zenith_no, azimuth_no))

    fig, ax = plt.subplots(subplot_kw=dict(projection='polar'), figsize=figsize)
    ax.pcolormesh(azimuths * np.pi / 180.0, zeniths, values, cmap=head_direction_input_colormap)
    # ax.set_yticks([])
    ax.set_thetagrids([0, 90, 180, 270])
    ax.tick_params(pad=-2)
    ax.set_ylim(0, 1)
    ax.grid(True)

    y_tick_labels = []
    ax.set_yticklabels(y_tick_labels)

    gridlines = ax.yaxis.get_gridlines()
    [line.set_linewidth(0.0) for line in gridlines]

    gridlines = ax.xaxis.get_gridlines()
    [line.set_linewidth(0.0) for line in gridlines]
    ax.axes.spines["polar"].set_visible(False)
    plt.subplots_adjust(left=0.25, right=0.75, bottom=0.25, top=0.75)
    # plt.show()
    if figname is not None:
        plt.savefig(figname, transparent=True)
        plt.close(fig)


if __name__ == "__main__":
    traj = Trajectory(TRAJ_NAME, add_time=False, dynamic_imports=Brian2MonitorResult)
    NO_LOADING = 0
    FULL_LOAD = 2
    traj.f_load(filename=os.path.join(DATA_FOLDER, TRAJ_NAME + ".hdf5"), load_parameters=FULL_LOAD,
                load_results=NO_LOADING)
    traj.v_auto_load = True
    save_figs = True

    print("# Plotting script polarized interneurons")
    print()

    map_length_scale = 200.0
    map_seed = 1
    exemplary_head_direction = 0

    print("## Map specifications")
    print("\tinput map scale: {:.1f} um".format(map_length_scale))
    print("\tmap seed: {:d}".format(map_seed))
    print()

    print("## Input specification")
    print("\tselected head direction: {:.0f}°".format(exemplary_head_direction))
    print()

    print("## Selected simulations")
    plot_scale = get_closest_scale(traj, map_length_scale)
    par_dict = {'seed': map_seed, 'scale': plot_scale}

    plot_run_names = filter_run_names_by_par_dict(traj, par_dict)
    run_name_dict = {}
    for run_name in plot_run_names:
        traj.f_set_crun(run_name)
        run_name_dict[traj.derived_parameters.runs[run_name].morphology.morph_label] = run_name
    for network_type, run_name in run_name_dict.items():
        print("{:s}: {:s}".format(network_type, run_name))

    directions = get_input_head_directions(traj)
    direction_idx = np.argmin(np.abs(np.array(directions) - np.deg2rad(exemplary_head_direction)))

    selected_neuron_excitatory = 1052
    selected_inhibitory_neuron = 28

    print("## Figure specification")
    print("\tpanel size: {:.2f} cm".format(panel_size * cm_per_inch))
    print()

    plot_colorbar(figsize=(0.8 * panel_size, 0.8 * panel_size), figname=FIGURE_SAVE_PATH + "A_i_colormap.svg")

    plot_input_map(traj, run_name_dict[POLARIZED], figname="A_i_exemplary_input_map" + FIGURE_SAVE_FORMAT,
                   figsize=(panel_size, panel_size))

    plot_example_input_maps(traj, figsize=(2 * panel_size, 2 * panel_size))

    plot_axonal_clouds(traj, plot_run_names)

    plot_firing_rate_map_excitatory(traj, direction_idx, plot_run_names, selected_neuron_excitatory)

    in_max_rate = plot_firing_rate_map_inhibitory(traj, direction_idx, plot_run_names, selected_inhibitory_neuron)

    #
    hdi_means = plot_hdi_histogram_combined_and_overlayed(
        traj, plot_run_names,
        selected_neuron_excitatory,
        selected_inhibitory_neuron,
        cut_off_dist=100.)
    #
    plot_polar_plot_excitatory(traj, plot_run_names, selected_neuron_excitatory)

    plot_polar_plot_inhibitory(traj, plot_run_names, selected_inhibitory_neuron)

    plot_firing_rate_similar_vs_diff_tuning(traj, plot_run_names, figsize=(1.2*panel_size, 1.2*panel_size))

    plot_exc_and_inh_hdi_over_simplex_grid_scale(traj, traj.f_get_run_names(), cut_off_dist=100.)

    plot_exc_and_inh_hdi_over_fit_corr_len(traj, traj.f_get_run_names(), cut_off_dist=100.)

    if not save_figs:
        plt.show()

    traj.f_restore_default()