structured_inhibition_spatial_network_model/scripts/spatial_network/perlin/paper_figures_spatial_head_direction_network_perlin_map.py

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.interneuron_placement import get_position_mesh, Pickle, get_correct_position_mesh
from scripts.model_figure.plot_circular_colorbar import plot_colorbar
from scripts.spatial_maps.simplex_input_tuning_correlation.correlation_length_distance_perlin import \
    correlation_length_fit_dict
from scripts.spatial_network.figures_spatial_head_direction_network_orientation_map import plot_hdi_in_space
from scripts.spatial_network.perlin.figure_utils import remove_frame, remove_ticks, add_length_scale, cm_per_inch, \
    panel_size, head_direction_input_colormap
from scripts.spatial_network.perlin.run_entropy_maximisation_perlin_map import DATA_FOLDER, TRAJ_NAME, \
    get_input_head_directions, POLARIZED, CIRCULAR, NO_SYNAPSES

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

FIGURE_SAVE_PATH = '../../../figures/figure_4_paper_perlin/svg_figures/'
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_correlation_length(traj, correlation_length):
    available_lengths = sorted(list(set(traj.f_get("correlation_length").f_get_range())))
    closest_length = available_lengths[np.argmin(np.abs(np.array(available_lengths) - correlation_length))]
    if closest_length != correlation_length:
        print("Warning: desired correlation length {:.1f} not available. Taking {:.1f} instead".format(
            correlation_length, 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', 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_correct_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_correct_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 plot_hdi_over_tuning(traj, plot_run_names):
    fig, ax = plt.subplots(1, 1)
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

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

        ex_tunings_plt = np.array(ex_tunings)
        sort_ids = ex_tunings_plt.argsort()
        ex_tunings_plt = ex_tunings_plt[sort_ids]

        head_direction_indices = traj.results[run_name].head_direction_indices

        hdi_plt = head_direction_indices
        hdi_plt = hdi_plt[sort_ids]
        ax.scatter(ex_tunings_plt, hdi_plt, label=label, alpha=0.3)

    ax.legend()
    ax.set_xlabel("Angles (rad)")
    ax.set_ylabel("head direction index")
    ax.set_title('hdi over input tuning', fontsize=16)

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


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_correct_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', 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, 'correlation_length': corresp_scale}

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

        X, Y = get_correct_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 = Pickle(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 = Pickle(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_correct_position_mesh(traj.results.runs[run_name].ex_positions)

        inhibitory_axonal_cloud_array = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        axonal_clouds = [Pickle(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_correlation_length(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_orientation_maps_diff_scales_with_ellipse(traj):
    n_ex = int(np.sqrt(traj.N_E))

    scale_run_names = []
    plot_scales = [0.0, 400.0, 800.0]
    real_scales = [get_closest_correlation_length(traj, scale) for scale in plot_scales]
    for scale in plot_scales:
        par_dict = {'seed': 1, 'correlation_length': get_closest_correlation_length(traj, scale), 'long_axis': 100.}
        scale_run_names.append(*filter_run_names_by_par_dict(traj, par_dict))
    print(scale_run_names)

    inhibitory_positions = traj.results.runs[scale_run_names[1]].inhibitory_axonal_cloud_array[:, :2]
    selected_polar_neuron = get_neurons_close_to_given_position((300, 300), 1, inhibitory_positions)[0]
    print(selected_polar_neuron)
    selected_circular_neuron = get_neurons_close_to_given_position((600, 600), 1, inhibitory_positions)[0]

    width = panel_size
    height = 1.5 * panel_size
    fig = plt.figure(figsize=(width, height))

    axes = ImageGrid(fig, 111, axes_pad=0.15,
                     nrows_ncols=(len(plot_scales), 1))
    for ax, run_name, scale in zip(axes, scale_run_names, real_scales):
        traj.f_set_crun(run_name)

        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 = [Pickle(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))
        # TODO: Why was this transposed for plotting? (now changed)
        c = ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap=head_direction_input_colormap)
        # ax.set_title('Correlation length: {}'.format(scale))
        # fig.colorbar(c, ax=ax, label="Tuning")
        ax.set_xticks([])
        ax.set_yticks([])

        p1 = axonal_clouds[selected_polar_neuron]
        ell = p1.get_ellipse()
        ell._linewidth = 2.
        ax.add_artist(ell)

        p2 = axonal_clouds[selected_circular_neuron]
        circ_r = 2 * np.sqrt(p2.a * p2.b)
        circ = Ellipse((p2.x, p2.y), circ_r, circ_r, fill=False, zorder=2, edgecolor='k')
        circ._linewidth = 2.

        ax.add_artist(circ)
        ax.annotate("{:.0f}".format(scale), xy=(1.0, 0.5), xytext=(2, 0), xycoords="axes fraction", textcoords="offset "
                                                                                                               "points",
                    va="center", ha="left")
        remove_frame(ax)

    # fig.suptitle('axonal cloud', fontsize=16)
    traj.f_restore_default()
    add_length_scale(axes[1], 200, -300, -100, 50, 50)
    axes[0].annotate("input maps", xy=(-0.5, 1.0), xycoords="axes fraction", rotation=90, ha="left", va="top")
    fig.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'F_orientation_maps_diff_scales_with_ellipse' + 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 = traj.derived_parameters.runs[run_name].morphology.morph_label
        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} {:.2f}'.format(short_labels(label), hdi)
        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])
    # ax.set_title('Firing Rate')
    # ax.plot([0.0, 0.0], [0.0, 1.05 * max_rate], color='red', alpha=0.25, linewidth=4.)
    # TODO: Set ticks for polar
    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]):
        # ax = axes[max_hdi_idx, run_idx]
        # label = traj.derived_parameters.runs[run_name].morphology.morph_label
        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} {:.2f}'.format(short_labels(label), hdi)
        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])
    # ax.set_title('Inh. Firing Rate')
    # TODO: Set ticks for polar
    # ticks = [np.round(max_rate / 3.), np.round(max_rate * 2. / 3.), np.round(max_rate)]
    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_over_corr_len(traj, plot_run_names):
    corr_len_expl = traj.f_get('correlation_length').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)

    hdi_frame = pd.Series(index=[corr_len_expl, seed_expl, label_expl])
    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
        head_direction_indices = traj.results[run_name].head_direction_indices
        hdi_frame[corr_len, seed, label] = np.mean(head_direction_indices)

    # TODO: Standart deviation also for the population
    hdi_exc_n_and_seed_mean = hdi_frame.groupby(level=[0, 2]).mean()
    hdi_exc_n_and_seed_std_dev = hdi_frame.groupby(level=[0, 2]).std()

    # Ellipsoid markers
    rx, ry = 5., 12.
    # area = rx * ry * np.pi * 2.
    area = 1.
    theta = np.arange(0, 2 * np.pi + 0.01, 0.1)
    verts = np.column_stack([rx / area * np.cos(theta), ry / area * np.sin(theta)])

    style_dict = {
        NO_SYNAPSES: ['grey', 'dashed', '', 0],
        POLARIZED: ['blue', 'solid', verts, 10.],
        CIRCULAR: ['lightblue', 'solid', 'o', 8.]

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

    fig, ax = plt.subplots(1, 1)

    for label in label_range:
        hdi_mean = hdi_exc_n_and_seed_mean[:, label]
        hdi_std = hdi_exc_n_and_seed_std_dev[:, label]
        corr_len_range = hdi_mean.keys().to_numpy()

        col, lin, mar, mar_size = style_dict[label]

        ax.plot(corr_len_range, hdi_mean, label=label, marker=mar, color=col, linestyle=lin, markersize=mar_size)
        plt.fill_between(corr_len_range, hdi_mean - hdi_std,
                         hdi_mean + hdi_std, alpha=0.4, color=col)
    ax.set_xlabel('Correlation length')
    ax.set_ylabel('Head Direction Index')
    ax.axvline(206.9, color='k', linewidth=0.5)
    ax.set_ylim(0.0, 1.0)
    ax.set_xlim(0.0, 400.)
    ax.legend()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'hdi_over_corr_len_scaled' + FIGURE_SAVE_FORMAT, dpi=200)
        plt.close(fig)


def plot_hdi_histogram_excitatory(traj, plot_run_names):
    labels = []
    hdis = []
    colors = ['black', 'red', 'green']
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        labels.append(label)

        head_direction_indices = traj.results.runs[run_name].head_direction_indices
        hdis.append(head_direction_indices)

    fig, ax = plt.subplots(1, 1, figsize=(6, 3))
    ax.hist(hdis, color=colors, label=labels, bins=30)

    for hdi, color in zip(hdis, colors):
        mean_hdi = np.mean(hdi)
        ax.axvline(mean_hdi, 0, 1, color=color, linestyle='--')
    ax.set_xlabel("HDI")
    ax.legend()

    fig.tight_layout()

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


def plot_hdi_violin_excitatory(traj, plot_run_names):
    labels = []
    hdis = []
    colors = ['black', 'red', 'green']
    no_conn_hdi = 0.
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        head_direction_indices = traj.results.runs[run_name].head_direction_indices

        if label == NO_SYNAPSES:
            no_conn_hdi = np.mean(head_direction_indices)
        else:
            labels.append(label)
            hdis.append(sorted(head_direction_indices))

    fig, ax = plt.subplots(1, 1, figsize=(6, 3))
    # hdis = np.array(hdis)
    viol_plt = ax.violinplot(hdis, showmeans=True, showextrema=False)

    viol_plt['cmeans'].set_color('black')

    for pc in viol_plt['bodies']:
        pc.set_facecolor('red')
        pc.set_edgecolor('black')
        pc.set_alpha(0.7)

    ax.axhline(no_conn_hdi, color='black', linestyle='--')
    ax.annotate(NO_SYNAPSES, xy=(0.45, 0.48), xycoords='axes fraction')

    ax.set_xticks(np.arange(1, len(labels) + 1))
    ax.set_xticklabels(labels)
    ax.set_ylabel('HDI')

    fig.tight_layout()

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


def plot_hdi_violin_inhibitory(traj, plot_run_names):
    labels = []
    hdis = []
    colors = ['black', 'red']
    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(label)

            head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
            hdis.append(sorted(head_direction_indices))

    fig, ax = plt.subplots(1, 1, figsize=(6, 3))
    viol_plt = ax.violinplot(hdis, showmeans=True, showextrema=False)

    viol_plt['cmeans'].set_color('black')

    for pc in viol_plt['bodies']:
        pc.set_facecolor('blue')
        pc.set_edgecolor('black')
        pc.set_alpha(0.7)

    ax.set_xticks(np.arange(1, len(labels) + 1))
    ax.set_xticklabels(labels)
    ax.set_ylabel('HDI')

    fig.tight_layout()

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


def plot_hdi_violin_combined(traj, plot_run_names):
    labels = []
    inh_hdis = []
    exc_hdis = []
    no_conn_hdi = 0.

    colors = ['black', 'red']
    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(label)

            inh_head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
            inh_hdis.append(sorted(inh_head_direction_indices))

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

    fig, ax = plt.subplots(1, 1, figsize=(6, 3))
    inh_viol_plt = ax.violinplot(inh_hdis, showmeans=True, showextrema=False)

    # viol_plt['cmeans'].set_color('black')
    #
    # for pc in viol_plt['bodies']:
    #     pc.set_facecolor('blue')
    #     pc.set_edgecolor('black')
    #     pc.set_alpha(0.7)

    for b in inh_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], m, np.inf)
        b.set_color('b')

    exc_viol_plt = ax.violinplot(exc_hdis, showmeans=True, showextrema=False)

    for b in exc_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], -np.inf, m)
        b.set_color('r')

    ax.axhline(no_conn_hdi, color='black', linestyle='--')
    ax.annotate(NO_SYNAPSES, xy=(0.45, 0.48), xycoords='axes fraction')

    ax.set_xticks(np.arange(1, len(labels) + 1))
    ax.set_xticklabels(labels)
    ax.set_ylabel('HDI')

    fig.tight_layout()

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


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

    in_polar_plot_hdi = []
    ex_polar_plot_hdi = []

    colors = ['black', 'red']
    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(label)

            inh_head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
            inh_hdis.append(sorted(inh_head_direction_indices))
            in_polar_plot_hdi.append(inh_head_direction_indices[in_polar_plot_id])

            exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
            exc_hdis.append(sorted(exc_head_direction_indices))
            ex_polar_plot_hdi.append(exc_head_direction_indices[ex_polar_plot_id])
        else:
            exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
            no_conn_hdi = np.mean(exc_head_direction_indices)
            ex_polar_plot_hdi.append(exc_head_direction_indices[ex_polar_plot_id])

    fig, ax = plt.subplots(1, 1, figsize=(3.5, 4.5))

    inh_ell_viol_plt = ax.violinplot(inh_hdis[0], showmeans=True, showextrema=False)
    for b in inh_ell_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], m, np.inf)
        b.set_color('b')
    mean_line = inh_ell_viol_plt['cmeans']
    mean_line.set_color('b')
    mean_line.get_paths()[0].vertices[:, 0] = np.clip(mean_line.get_paths()[0].vertices[:, 0], m, np.inf)

    exc_ell_viol_plt = ax.violinplot(exc_hdis[0], showmeans=True, showextrema=False)
    for b in exc_ell_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], m, np.inf)
        b.set_color('r')
    mean_line = exc_ell_viol_plt['cmeans']
    mean_line.set_color('r')
    mean_line.get_paths()[0].vertices[:, 0] = np.clip(mean_line.get_paths()[0].vertices[:, 0], m, np.inf)

    inh_cir_viol_plt = ax.violinplot(inh_hdis[1], showmeans=True, showextrema=False)
    for b in inh_cir_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], -np.inf, m)
        b.set_color('b')
    mean_line = inh_cir_viol_plt['cmeans']
    mean_line.set_color('b')
    mean_line.get_paths()[0].vertices[:, 0] = np.clip(mean_line.get_paths()[0].vertices[:, 0], -np.inf, m)

    exc_cir_viol_plt = ax.violinplot(exc_hdis[1], showmeans=True, showextrema=False)
    for b in exc_cir_viol_plt['bodies']:
        m = np.mean(b.get_paths()[0].vertices[:, 0])
        b.get_paths()[0].vertices[:, 0] = np.clip(b.get_paths()[0].vertices[:, 0], -np.inf, m)
        b.set_color('r')
    mean_line = exc_cir_viol_plt['cmeans']
    mean_line.set_color('r')
    mean_line.get_paths()[0].vertices[:, 0] = np.clip(mean_line.get_paths()[0].vertices[:, 0], -np.inf, m)

    ax.axhline(no_conn_hdi, 0.5, 1., color='black', linestyle='--')
    ax.axvline(1.0, color='k')
    ax.annotate(NO_SYNAPSES, xy=(0.75, 0.415), xycoords='axes fraction')
    ax.set_xlim(0.5, 1.5)
    ax.set_ylim(0.0, 1.0)
    ax.set_xticks([0.75, 1.25])
    ax.set_xticklabels([CIRCULAR, POLARIZED])
    ax.set_ylabel('HDI')

    fig.tight_layout()

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

    return ex_polar_plot_hdi, in_polar_plot_hdi


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.map.sheet_size - cut_off_dist) & \
                              (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                              (inh_axonal_cloud[:, 1] <= traj.parameters.map.sheet_size - cut_off_dist)
            # print(inh_positions)
            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
            # print(ex_positions)
            exc_cut_off_ids = (ex_positions[:, 0] >= cut_off_dist) & \
                              (ex_positions[:, 0] <= traj.parameters.map.sheet_size - cut_off_dist) & \
                              (ex_positions[:, 1] >= cut_off_dist) & \
                              (ex_positions[:, 1] <= traj.parameters.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]

    # # fig = plt.figure(figsize=(3.5, 3.5))
    # # gs1 = gridspec.GridSpec(1, 2)
    #
    # fig = plt.figure(constrained_layout=True)
    # gs = gridspec.GridSpec(ncols=1, nrows=2, hspace=0.0, wspace=0.0, figure=fig)
    # # gs.update(wspace=0.0, hspace=0.0)  # set the spacing between axes.
    # # f2_ax1 = fig.add_subplot(gs[0, 0])
    # # f2_ax2 = fig.add_subplot(gs[0, 1])
    # print(gs.get_subplot_params())

    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_cutoff_{}um' + 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 plot_hdi_histogram_inhibitory(traj, plot_run_names, in_polar_plot_id):
    labels = []
    hdis = []
    colors = ['black', 'red']
    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(label)

            head_direction_indices = traj.results.runs[run_name].inh_head_direction_indices
            print('inh {}: {}'.format(label, head_direction_indices[in_polar_plot_id]))
            hdis.append(head_direction_indices)

    fig, ax = plt.subplots(1, 1, figsize=(6, 3))
    ax.hist(hdis, color=colors, label=labels, bins=30)

    for hdi, color in zip(hdis, colors):
        mean_hdi = np.mean(hdi)
        ax.axvline(mean_hdi, 0, 1, color=color, linestyle='--')
    ax.set_xlabel("HDI")
    ax.legend()

    fig.tight_layout()

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


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('correlation_length').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.map.sheet_size - cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] <= traj.parameters.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.map.sheet_size - cut_off_dist) & \
                          (ex_positions[:, 1] >= cut_off_dist) & \
                          (ex_positions[:, 1] <= traj.parameters.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, 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[1, 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()
        print(label)
        for corr_len, ex_hdi, in_hdi in zip(corr_len_range, exc_hdi_mean, inh_hdi_mean):
            print("length: {:.2f} um, ex hdi: {:.2f} and in hdi {:.2f}".format(corr_len, ex_hdi, in_hdi))

        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_correlation_length(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('correlation_length').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.map.sheet_size - cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] >= cut_off_dist) & \
                          (inh_axonal_cloud[:, 1] <= traj.parameters.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.map.sheet_size - cut_off_dist) & \
                          (ex_positions[:, 1] >= cut_off_dist) & \
                          (ex_positions[:, 1] <= traj.parameters.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', 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[1, 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()
        print(label)
        for corr_len, ex_hdi, in_hdi in zip(corr_len_range, exc_hdi_mean, inh_hdi_mean):
            print("length: {:.2f} um, ex hdi: {:.2f} and in hdi {:.2f}".format(corr_len, ex_hdi, in_hdi))

        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')
    print('The used correlation length is ', corr_len_fit_dict[get_closest_correlation_length(traj, 200.0)])
    ax.axvline(corr_len_fit_dict[get_closest_correlation_length(traj, 200.0)], 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_corr_len_scaled' + FIGURE_SAVE_FORMAT)
        plt.close(fig)

def plot_in_degree_map(traj, plot_run_names):
    n_ex = int(np.sqrt(traj.N_E))
    max_degree = 0
    for run_name in plot_run_names:
        ie_adjacency = traj.results.runs[run_name].ie_adjacency
        exc_degree = np.sum(ie_adjacency, axis=0)
        run_max_degree = np.max(exc_degree)
        if run_max_degree > max_degree:
            max_degree = run_max_degree

    fig, axes = plt.subplots(1, 2, figsize=(9., 4.5))
    for ax, run_name in zip(axes, plot_run_names[:-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)

        number_of_excitatory_neurons_per_row = int(np.sqrt(traj.N_E))
        ie_adjacency = traj.results.runs[run_name].ie_adjacency
        exc_degree = np.sum(ie_adjacency, axis=0)
        c = ax.pcolor(X, Y, np.reshape(exc_degree, (number_of_excitatory_neurons_per_row,
                                                    number_of_excitatory_neurons_per_row)), vmin=0, vmax=max_degree,
                      cmap='hot')
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="in/out-degree")
    fig.suptitle('in/out-degree', fontsize=16)
    traj.f_restore_default()

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


def plot_spatial_hdi_map(traj, plot_run_names):
    max_val = 0
    for run_name in plot_run_names:
        hdis = traj.results.runs[run_name].head_direction_indices
        run_max_val = np.max(hdis)
        if run_max_val > max_val:
            max_val = run_max_val

    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.5))
    for ax, run_name in zip(axes, plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        positions = traj.results.runs[run_name].ex_positions
        head_direction_indices = traj.results[run_name].head_direction_indices
        print('Mean {}-HDI = {}'.format(label, np.mean(head_direction_indices)))
        c = plot_hdi_in_space(ax, positions, head_direction_indices, max_val)
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="HDI")
    fig.suptitle('spatial HDI map', fontsize=16)

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


def plot_exc_spatial_hdi_map(traj, plot_run_names):
    max_val = 0
    for run_name in plot_run_names:
        hdis = traj.results.runs[run_name].head_direction_indices
        run_max_val = np.max(hdis)
        if run_max_val > max_val:
            max_val = run_max_val

    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.5))
    for ax, run_name in zip(axes, plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        positions = traj.results.runs[run_name].ex_positions
        head_direction_indices = traj.results[run_name].head_direction_indices
        print('Mean {}-HDI = {}'.format(label, np.mean(head_direction_indices)))
        c = plot_hdi_in_space(ax, positions, head_direction_indices, max_val)
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="HDI")
    fig.suptitle('spatial exc. HDI map', fontsize=16)

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


def plot_inh_spatial_hdi_map(traj, plot_run_names):
    max_val = 0
    for run_name in plot_run_names:
        hdis = traj.results.runs[run_name].inh_head_direction_indices
        run_max_val = np.max(hdis)
        if run_max_val > max_val:
            max_val = run_max_val

    fig, axes = plt.subplots(1, 3, figsize=(13.5, 4.5))
    for ax, run_name in zip(axes, plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        ax_cloud = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        positions = [[x, y] for x, y, phi in ax_cloud]
        head_direction_indices = traj.results[run_name].inh_head_direction_indices
        print('Mean {}-HDI = {}'.format(label, np.mean(head_direction_indices)))
        c = plot_hdi_in_space(ax, positions, head_direction_indices, max_val)
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="HDI")
    fig.suptitle('spatial inh. HDI map', fontsize=16)

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'spatial_inh_hdi_map' + FIGURE_SAVE_FORMAT, dpi=200)
        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)
    print(dir_bins)

    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'$\theta_{pref} \pm 30°$')
    # rects2 = ax.bar(x + width / 2, plot_fr_array[:, 1], width,
    #                 label=r'$\theta_{opp} \pm 30°$')
    rects1 = ax.bar(x - width / 2, plot_fr_array[:, 0], width,
                    label=r'sim.')
    rects2 = ax.bar(x + width / 2, plot_fr_array[:, 1], width,
                    label=r'diff.')

    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 + '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]

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

    # corr_len_fit_dict = correlation_length_fit_dict(traj, map_type='perlin', load=True)
    # plt.plot(corr_len_fit_dict.keys(),corr_len_fit_dict.values())
    # plt.show()
    # abbrechen


    print("## Map specifications")
    print("\tcorrelation length: {:.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_corr_len = get_closest_correlation_length(traj, map_length_scale)
    par_dict = {'seed': map_seed, 'correlation_length': plot_corr_len}

    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.)
    #
    number_of_suggestions = 20
    representative_excitatory_neuron_indices = get_neurons_with_given_hdi(hdi_means["polar_exc"], hdi_means[
        "circular_exc"],
                                                                          number_of_suggestions, plot_run_names,
                                                                          traj, "ex")
    representative_inhibitory_neuron_indices = get_neurons_with_given_hdi(hdi_means["polar_inh"],
                                                                          hdi_means["circular_inh"],
                                                                          number_of_suggestions, plot_run_names,
                                                                          traj, "in")

    plot_polar_plot_excitatory(traj, plot_run_names, selected_neuron_excitatory)

    for suggested_excitatory_neuron in representative_excitatory_neuron_indices:
        preferred_polar, opposite_polar = get_firing_rates_along_preferred_axis(
            traj, run_name_dict["ellipsoid"], suggested_excitatory_neuron)
        preferred_circular, opposite_circular = get_firing_rates_along_preferred_axis(
                    traj, run_name_dict["circular"], suggested_excitatory_neuron)

        plot_polar_plot_excitatory(traj, plot_run_names, suggested_excitatory_neuron,
                                   figname=FIGURE_SAVE_PATH + "X_polar_plot_excitatory_neuron_signal_increase_{"
                                                              ":0>2.0f}_noise_decrease_{"
                                                              ":0>2.0f}_neuron_id_{"
                                                              ":d}.png".format(
                                       preferred_polar-preferred_circular,
                                       opposite_circular-opposite_polar, suggested_excitatory_neuron))


    plot_polar_plot_inhibitory(traj, plot_run_names, selected_inhibitory_neuron)

    for suggested_inhibitory_neuron in representative_inhibitory_neuron_indices:
        polar_hdi = get_hdi(traj, run_name_dict["ellipsoid"], suggested_inhibitory_neuron, "in")
        circular_hdi = get_hdi(traj, run_name_dict["circular"], suggested_inhibitory_neuron, "in")
        plot_polar_plot_inhibitory(traj, plot_run_names, suggested_inhibitory_neuron,
                                   figname=FIGURE_SAVE_PATH + "X_polar_plot_inhibitory_neuron_polarized_{"
                                                              ":.2f}_circular_{:.2f}_neuron_id_{"
                                                              ":d}.png".format(polar_hdi, circular_hdi,
                                       suggested_inhibitory_neuron))

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

    plot_orientation_maps_diff_scales_with_ellipse(traj)

    # 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()