import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from brian2.units import *
from matplotlib import gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
from pypet import Trajectory
from pypet.brian2 import Brian2MonitorResult
from scipy.optimize import curve_fit
from matplotlib.patches import Ellipse
from matplotlib.patches import Polygon
import matplotlib.legend as mlegend
from matplotlib.patches import Rectangle


from scripts.interneuron_placement import get_position_mesh, Pickle, get_correct_position_mesh
from scripts.spatial_network.figures_spatial_head_direction_network_orientation_map import plot_hdi_in_space
from scripts.spatial_network.run_entropy_maximisation_orientation_map import DATA_FOLDER, TRAJ_NAME

# FIGURE_SAVE_PATH = '../../figures/figure_4_paper_draft/'
FIGURE_SAVE_PATH = '../../figures/figure_4_poster_fens/'


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 gauss(x, *p):
    A, mu, sigma, B = p
    return A * np.exp(-(x - mu) ** 2 / (2. * sigma ** 2)) + B


def plot_tuning_curve(traj, direction_idx, plot_run_names):
    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)

    rate_frame = pd.Series(index=[seed_expl, label_expl])
    rate_frame.index.names = ["seed", "label"]

    dir_bins = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    rate_bins = [[] for i in range(len(dir_bins)-1)]

    for run_name, seed, label in zip(plot_run_names, seed_expl, label_expl):
        ex_tunings = traj.results.runs[run_name].ex_tunings
        binned_idx = np.digitize(ex_tunings, dir_bins)

        #TODO: Avareage over directions by recentering
        firing_rate_array = traj.results[run_name].firing_rate_array
        for bin_idx, rate in zip(binned_idx, firing_rate_array[:, direction_idx]):
            rate_bins[bin_idx].append(rate)

        rate_bins_mean = [np.mean(rate_bin) for rate_bin in rate_bins]

        rate_frame[seed, label] = firing_rate_array

    # TODO: Standart deviation also for the population
    rate_seed_mean = rate_frame.groupby(level=[1]).mean()
    rate_seed_std_dev = rate_frame.groupby(level=[1]).std()


    style_dict = {
        'no conn': ['grey', 'dashed', '', 0],
        'ellipsoid': ['blue', 'solid', 'x', 10.],
        'circular': ['lightblue', 'solid', 'o', 8.]

    }

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

    for label in label_range:
        hdi_mean = rate_seed_mean[label]
        hdi_std = rate_seed_std_dev[label]
        ex_tunings = traj.results.runs[run_name].ex_tunings

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


    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

        coeff_list = []

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

        firing_rate_array = traj.results[run_name].firing_rate_array
        # firing_rate_array = traj.f_get('firing_rate_array')

        rates_plt = firing_rate_array[:, direction_idx]
        rates_plt = rates_plt[sort_ids]
        ax.scatter(ex_tunings_plt, rates_plt / hertz, label=label, alpha=0.3)

    ax.legend()
    ax.set_xlabel("Angles (rad)")
    ax.set_ylabel("f (Hz)")
    ax.set_title('tuning curves', fontsize=16)

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'tuning_curve.png', dpi=200)

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):
    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 == 'ellipsoid':
            #     n_id_max_rate = np.argmax(f_rates)
            max_val = run_max_val

    n_id_polar_plot = 609

    # 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 = 13.
    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 = 3.5
    # color_bar_size = 0.05 * height + 0.05
    # width = 3 * height + color_bar_size
    width = 10.5

    fig, axes = plt.subplots(1, 3, figsize=(width, height))
    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(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'))
    # fig.suptitle('spatial firing rate map', fontsize=16)
    colorbar(c)


    fig.tight_layout()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'firing_rate_map.png', dpi=200)

    return n_id_polar_plot

def plot_firing_rate_map_inhibitory(traj, direction_idx, plot_run_names):
    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 = 52

    # 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]

    fig, axes = plt.subplots(1, 2, figsize=(7.0, 3.5))
    for ax, run_name in zip(axes, [plot_run_names[i] for i in [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(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))
        # fig.colorbar(c, ax=ax, label="f (Hz)")
    # fig.suptitle('spatial firing rate map', fontsize=16)
    colorbar(c)
    fig.tight_layout()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'inh_firing_rate_map.png', dpi=200)
    return n_id_polar_plot, 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.png', dpi=200)


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

    fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.5))
    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))
        # TODO: Why was this transposed for plotting? (now changed)
        c = ax.pcolor(X, Y, head_dir_preference, vmin=-np.pi, vmax=np.pi, cmap='hsv')
        ax.set_title(label)
        # fig.colorbar(c, ax=ax, label="Tuning")

        if label != 'no conn' and axonal_clouds is not None:
            for i, p in enumerate(axonal_clouds):
                ell = p.get_ellipse()
                ax.add_artist(ell)
    # fig.suptitle('axonal cloud', fontsize=16)
    traj.f_restore_default()

    fig.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'axonal_clouds.png', dpi=200)

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.png', dpi=200)


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

    scale_run_names = []
    plot_scales = [0.0, 100.0, 200.0, 300.0, 400.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))
    print(scale_run_names)

    fig, axes = plt.subplots(1, 5, 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)

        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='hsv')
        # ax.set_title('Correlation length: {}'.format(scale))
        # fig.colorbar(c, ax=ax, label="Tuning")
        ax.set_xticks([])
        ax.set_yticks([])

        p1 = axonal_clouds[44]
        ell = p1.get_ellipse()
        ell._linewidth = 5.
        ax.add_artist(ell)

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

        ax.add_artist(circ)

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

    fig.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'orientation_maps_diff_scales_with_ellipse.png', dpi=200)

def plot_excitatory_condensed_polar_plot(traj, plot_run_names, polar_plot_id, ex_polar_plot_hdi):
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    directions_plt = list(directions)
    directions_plt.append(directions[0])
    fig, ax = plt.subplots(1, 1, figsize=(3.5, 3.5), 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 = polar_plot_id

    line_styles = ['dotted', 'solid', 'dashed']
    colors = ['r', 'lightsalmon', 'grey']

    max_rate = 0.0
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        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])
        ax.plot(directions_plt, rate_plot, label='{0}, HDI: {1:.2f}'.format(label, ex_polar_plot_hdi[run_idx]), color=colors[run_idx], linestyle=line_styles[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 = [30., 60., 90.]
    ax.set_rticks(ticks)
    ax.set_rlabel_position(230)
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15),
          fancybox=True, shadow=True)
    plt.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'condensed_polar_plot.png', dpi=200)

def plot_inhibitory_condensed_polar_plot(traj, plot_run_names, polar_plot_id, in_polar_plot_hdi):
    directions = np.linspace(-np.pi, np.pi, traj.input.number_of_directions, endpoint=False)
    directions_plt = list(directions)
    directions_plt.append(directions[0])
    fig, ax = plt.subplots(1, 1, figsize=(3.5, 3.5), 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 = polar_plot_id

    line_styles = ['dotted', 'solid']
    colors = ['b', 'lightblue']

    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

        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])
        ax.plot(directions_plt, rate_plot, label='{0}, HDI: {1:.2f}'.format(label, in_polar_plot_hdi[run_idx]), color=colors[run_idx], linestyle=line_styles[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_rticks(ticks)
    ax.set_rlabel_position(230)
    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15),
              fancybox=True, shadow=True)
    plt.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'condensed_inhibitory_polar_plot.png', dpi=200)

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 conn': ['grey', 'dashed', '', 0],
        'ellipsoid': ['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.png', dpi=200)


def plot_hdi_histogram_excitatory(traj, plot_run_names, ex_polar_plot_id):
    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
        print('exc {}: {}'.format(label, head_direction_indices[ex_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_excitatory.png', dpi=200)


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 conn':
            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 conn', 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.png', dpi=200)


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 conn':
            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.png', dpi=200)

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 conn':
            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 conn', 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)

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 conn':
            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 conn', 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', 'ellipsoid'])
    ax.set_ylabel('HDI')


    fig.tight_layout()

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

    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):
    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 conn':
            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 = 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())

    fig, axes = plt.subplots(2, 1, figsize=(3.5, 2.5))
    plt.subplots_adjust(wspace=0, hspace=0)
    bins = np.linspace(0.0, 1.0, 21, endpoint=True)

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

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

        ax.axvline(np.mean(exc_hdis[i]), color='r')
        ax.axvline(np.mean(inh_hdis[i]), color='b')
        ax.axvline(no_conn_hdi, color='grey', linestyle='--')

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

    # ax.axhline(no_conn_hdi, 0.5, 1., color='black', linestyle='--')
    # ax.axvline(1.0, color='k')
    # ax.annotate('no conn', 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', 'ellipsoid'])
    # ax.set_ylabel('HDI')


    # fig.tight_layout()

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

    return ex_polar_plot_hdi, in_polar_plot_hdi

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 conn':
            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.png', dpi=200)


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_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)

    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
        head_direction_indices = traj.results[run_name].head_direction_indices
        #TODO: Actual correlation lengths
        # actual_corr_len = get_correlation_length(ex_tunings.reshape((30,30)), 450, 30)
        exc_hdi_frame[corr_len, seed, label] = np.mean(head_direction_indices)

        inh_head_direction_indices = traj.results[run_name].inh_head_direction_indices
        inh_hdi_frame[corr_len, seed, label] = np.mean(inh_head_direction_indices)

    # TODO: Standart 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()

    exc_style_dict = {
        'no conn': ['grey', 'dashed', '', 0],
        'ellipsoid': ['red', 'solid', '^', 8.],
        'circular': ['lightsalmon', 'solid', '^', 8.]
    }
    inh_style_dict = {
                     'no conn': ['grey', 'dashed', '', 0],
                     'ellipsoid': ['blue', 'solid', 'o', 8.],
                     'circular': ['lightblue', 'solid', 'o', 8.]
    }

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

    for label in label_range:
        if label == 'no conn':
            ax.axhline(exc_hdi_n_and_seed_mean[0, label], color='grey', linestyle='--')
            ax.annotate('input', xy=(1.01, 0.44), xycoords='axes fraction')

            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]


        ax.plot(corr_len_range, 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(corr_len_range, exc_hdi_mean - exc_hdi_std,
                         exc_hdi_mean + exc_hdi_std, alpha=0.3, color=exc_col)
        ax.plot(corr_len_range, 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(corr_len_range, inh_hdi_mean - inh_hdi_std,
                         inh_hdi_mean + inh_hdi_std, alpha=0.3, color=inh_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.)
    tablelegend(ax, ncol=2, bbox_to_anchor=(1, 1),
                row_labels=['exc.', 'inh.'],
                col_labels=['ellipsoid', 'circular'],
                title_label='')
    # plt.legend()
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'exc_and_inh_hdi_over_corr_len_scaled.png', dpi=200)

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.png', dpi=200)

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="head direction index")
    fig.suptitle('spatial HDI map', fontsize=16)

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

def plot_ellipse_hdi_over_circular_hdi(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 conn':
            labels.append(label)

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

            exc_head_direction_indices = traj.results.runs[run_name].head_direction_indices
            exc_hdis.append(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=(4.5, 4.5))
    plt.subplots_adjust(wspace=0, hspace=0)
    bins = np.linspace(0.0, 1.0, 21, endpoint=True)

    ax.scatter(exc_hdis[1], exc_hdis[0], color='r', alpha=0.3)
    ax.scatter(inh_hdis[1], inh_hdis[0], color='b', alpha=0.3)
    ax.set_xlabel('{} HDI'.format(labels[1]))
    ax.set_ylabel('{} HDI'.format(labels[0]))
    ax.plot([0.0, 1.0], [0.0, 1.0], 'k')

    # fig.tight_layout()

    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'ellipse_hdi_over_circular_hdi.png', dpi=200)

def plot_hdi_over_in_degree_binned(traj, plot_run_names):
    fig, ax = plt.subplots(1, 1, figsize=(4.5, 4.5))
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        if label == 'no conn':
            continue
        ie_adjacency = traj.results.runs[run_name].ie_adjacency
        exc_degree = np.sum(ie_adjacency, axis=0)

        head_direction_indices = traj.results[run_name].head_direction_indices


        n_bins = 13
        bins = np.arange(n_bins + 1) - 0.5
        exc_degree_bin_ids = np.digitize(exc_degree, bins)
        # print(bins)
        # print(exc_degree)
        # print(exc_degree_bin_ids)
        hdis_per_bin = [[] for i in range(n_bins)]

        for idx, hdi in zip(exc_degree_bin_ids, head_direction_indices):
            hdis_per_bin[idx - 1].append(hdi)

        mean_hdi_per_bin = [np.mean(hdis) for hdis in hdis_per_bin]

        ax.bar(np.arange(n_bins), mean_hdi_per_bin, label=label, alpha=0.3)

    ax.legend()
    ax.set_xlabel("In-degree of exc. neurons")
    ax.set_ylabel("Mean Head Direction Index")
    # ax.set_title('hdi over in-degree', fontsize=16)
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'hdi_over_in_degree_binned.png', dpi=200)

def plot_in_degree_hist(traj, plot_run_names):
    fig, ax = plt.subplots(1, 1, figsize=(4.5, 4.5))
    for run_idx, run_name in enumerate(plot_run_names):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label
        if label == 'no conn':
            continue
        ie_adjacency = traj.results.runs[run_name].ie_adjacency
        exc_degree = np.sum(ie_adjacency, axis=0)

        head_direction_indices = traj.results[run_name].head_direction_indices


        n_bins = 13
        bins = np.arange(n_bins + 1) - 0.5
        in_degree_hist, bin_edges = np.histogram(exc_degree, bins)
        print(in_degree_hist)

        # ax.hist(exc_degree, bins=bins, label=label, alpha=0.3)
        ax.bar(np.arange(n_bins), in_degree_hist, label=label, alpha=0.3)

    ax.legend()
    ax.set_xlabel("In-degree of exc. neurons")
    ax.set_ylabel("# exc. neurons")
    # ax.set_title('hdi over in-degree', fontsize=16)
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'in_degree_hist.png', dpi=200)

def plot_hdi_over_in_degree(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
        if label == 'no conn':
            continue
        ie_adjacency = traj.results.runs[run_name].ie_adjacency
        exc_degree = np.sum(ie_adjacency, axis=0)

        sort_ids = exc_degree.argsort()
        exc_degree_plt = exc_degree[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(exc_degree_plt, hdi_plt, label=label, alpha=0.3)

    ax.legend()
    ax.set_xlabel("In-degree of exc. neurons")
    ax.set_ylabel("Head Direction Index")
    ax.set_title('hdi over in-degree', fontsize=16)
    if save_figs:
        plt.savefig(FIGURE_SAVE_PATH + 'hdi_over_in_degree.png', dpi=200)

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

    plot_corr_len = get_closest_correlation_length(traj, 200.0)
    par_dict = {'seed': 1, 'correlation_length': plot_corr_len}
    plot_run_names = filter_run_names_by_par_dict(traj, par_dict)
    print(plot_run_names)

    direction_idx = 6
    dir_indices = [0, 3, 6, 9]

    # plot_axonal_clouds(traj, plot_run_names)
    # #
    ex_polar_plot_id = plot_firing_rate_map_excitatory(traj, direction_idx, plot_run_names)
    #
    in_polar_plot_id, in_max_rate = plot_firing_rate_map_inhibitory(traj, direction_idx, plot_run_names)
    #
    ex_polar_plot_hdi, in_polar_plot_hdi = plot_hdi_histogram_combined_and_overlayed(traj, plot_run_names, ex_polar_plot_id, in_polar_plot_id)

    plot_excitatory_condensed_polar_plot(traj, plot_run_names, ex_polar_plot_id, ex_polar_plot_hdi)

    plot_inhibitory_condensed_polar_plot(traj, plot_run_names, in_polar_plot_id, in_polar_plot_hdi)

    # plot_ellipse_hdi_over_circular_hdi(traj, plot_run_names)
    #
    # plot_hdi_over_in_degree_binned(traj, plot_run_names)
    #
    # plot_hdi_over_in_degree(traj, plot_run_names)
    #
    # plot_in_degree_hist(traj, plot_run_names)

    #
    # plot_orientation_maps_diff_scales_with_ellipse(traj)
    #
    # plot_hdi_histogram_inhibitory(traj, plot_run_names, in_polar_plot_id)
    #
    # plot_hdi_histogram_excitatory(traj, plot_run_names, ex_polar_plot_id)
    #
    # plot_hdi_over_corr_len(traj, traj.f_get_run_names())

    # plot_hdi_violin_combined(traj, plot_run_names)
    #

    # plot_exc_and_inh_hdi_over_corr_len(traj, traj.f_get_run_names())

    # plot_in_degree_map(traj, plot_run_names)
    #
    # plot_spatial_hdi_map(traj, plot_run_names)

    # par_dict = {'correlation_length': plot_corr_len}
    # single_corr_len_run_names = filter_run_names_by_par_dict(traj, par_dict)
    # plot_tuning_curve(traj, direction_idx, single_corr_len_run_names)

    if not save_figs:
        plt.show()

    traj.f_restore_default()