import numpy as np
from pypet import Trajectory
from pypet.brian2 import Brian2MonitorResult
from tqdm import tqdm
from brian2.units import *
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


from scripts.interneuron_placement import create_grid_of_excitatory_neurons, get_position_mesh, Pickle
from scripts.spatial_network.head_direction_index_over_noise_scale import calculate_rates
from scripts.spatial_network.run_entropy_maximisation_orientation_map import DATA_FOLDER, TRAJ_NAME, get_orientation_map

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

    # p0 is the initial guess for the fitting coefficients (A, mu and sigma above)
    p0 = [30., 0., 1., 10.]
    # run_name = traj.v_crun

    fig, ax = plt.subplots(1, 1)
    for run_idx, run_name in enumerate(traj.f_get_run_names()):
        label = traj.derived_parameters.runs[run_name].morphology.morph_label

        ex_tunings = traj.results.runs[run_name].ex_tunings
        # ex_positions = traj.f_get('ex_positions')

        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]
        coeff, var_matrix = curve_fit(gauss, ex_tunings_plt, rates_plt / hertz, p0=p0)
        coeff_list.append(coeff)
        print('{}-std_dev = {}'.format(label, np.abs(coeff[2])))
        ax.scatter(ex_tunings_plt, rates_plt / hertz, label=label, alpha=0.3)
        ax.plot(ex_tunings_plt, gauss(ex_tunings_plt, *coeff), label=label + '-fit')

    # ax.plot(ex_tunings,gauss(ex_tunings,*coeff_list[1]),'ro')
    # ax.plot(ex_tunings,gauss(ex_tunings,*coeff_list[2]),'bo')
    ax.legend()
    ax.set_xlabel("Angles (rad)")
    ax.set_ylabel("f (Hz)")
    ax.set_title('tuning curves', fontsize=16)

def plot_firing_rate_map(traj, direction_idx):

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

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)
        firing_rate_array = traj.results[run_name].firing_rate_array
        number_of_excitatory_neurons_per_row = int(np.sqrt(traj.N_E))
        c = ax.pcolor(X, Y, np.reshape(firing_rate_array[:, direction_idx], (number_of_excitatory_neurons_per_row,
                                                       number_of_excitatory_neurons_per_row)), vmin=0, cmap='hot')
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="f (Hz)")
    fig.suptitle('spatial firing rate map', fontsize=16)


def plot_spatial_hdi_map(traj):

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

        X, Y = get_position_mesh(traj.results.runs[run_name].ex_positions)
        head_direction_indices = traj.results[run_name].head_direction_indices
        number_of_excitatory_neurons_per_row = int(np.sqrt(traj.N_E))
        c = ax.pcolor(X, Y, np.reshape(head_direction_indices, (number_of_excitatory_neurons_per_row,
                                                       number_of_excitatory_neurons_per_row)), vmin=0, cmap='hot')
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="f (Hz)")
    fig.suptitle('spatial HDI map', fontsize=16)


def plot_neural_sheet(traj):

    n_ex = int(np.sqrt(traj.N_E))

    fig, axes = plt.subplots(1, 2, figsize=(13.5, 4.5))
    for ax, run_name in zip(axes, traj.f_get_run_names()):
        traj.f_set_crun(run_name)

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

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

        inhibitory_axonal_cloud_array = traj.results.runs[run_name].inhibitory_axonal_cloud_array
        print(run_name,traj.morphology.long_axis,traj.morphology.short_axis)
        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.T, vmin=-np.pi, vmax=np.pi, cmap='twilight')
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="Tuning")


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

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

    fig, axes = plt.subplots(1, 2, figsize=(13.5, 4.5))
    for ax, run_name in zip(axes, traj.f_get_run_names()):
        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)), cmap='twilight')
        ax.set_title(label)
        fig.colorbar(c, ax=ax, label="Tuning")
    fig.suptitle('in/out-degree', fontsize=16)


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

    direction_idx = 6

    plot_tuning_curve(traj, direction_idx)

    plot_firing_rate_map(traj, direction_idx)

    plot_spatial_hdi_map(traj)

    plot_neural_sheet(traj)

    plot_degree_on_sheet(traj)

    plt.show()

    traj.f_restore_default()
