import matplotlib.pyplot as plt
import noise
import numpy as np
from brian2.units import *
from scripts.spatial_maps.orientation_map import OrientationMap
from scripts.interneuron_placement import create_interneuron_sheet_entropy_max_orientation, \
    create_grid_of_excitatory_neurons, plot_neural_sheet, create_interneuron_sheet_by_repulsive_force, \
    get_excitatory_neurons_in_inhibitory_axonal_clouds, get_position_mesh

import scripts.models as modellib
from scripts.ring_network.head_direction import get_head_direction_input, \
    ex_in_network
from scripts.spatial_network.analysis_spatio_temporal_activity import transform_spike_train_in_activity_array, \
    create_neural_sheet_animation

save_dont_show = False
# output_folder = '../spatial_network_output/entropy_rotation'
output_folder = '../../spatial_network_output/'

'''
Neural sheet
'''

# To go from the 3d data obtained in Peng 2017 to a 2d sheet is not straightforward
# Assuming a depth of 100um, the density becomes a 1000 interneurons/mm2 and thus about 3000-4000 exneurons/mm2
# A long 200um short 50um interneuron axon would cover about 90 partners
# On the other hand, the typical contacts per excitatory neuron would be A_total/(A_I*N_I) = 30


N_E = 900
N_I = 90

sheet_x = 450 * um
sheet_y = 450 * um

inhibitory_axon_long_axis = 100 * um
inhibitory_axon_short_axis = 25 * um

number_of_excitatory_neurons_per_row = int(np.sqrt(N_E))

'''
Tuning Maps
'''

# tuning_label = "Perlin"
tuning_label = "Orientation map"

# optimization_label = "Repulsive"
optimization_label = "Entropy Optimization"

if tuning_label == "Perlin":
    tuning_map = lambda x, y: noise.pnoise2(x / 100.0, y / 100.0, octaves=2)*np.pi
elif tuning_label == "Orientation map":
    map = OrientationMap(number_of_excitatory_neurons_per_row + 1,number_of_excitatory_neurons_per_row + 1,
                         200,sheet_x/um,sheet_y/um,1)
    # map.improve(10)
    map.load_orientation_map()
    tuning_map = lambda x, y: map.orientation_map(x, y)


ex_positions, ex_tunings = create_grid_of_excitatory_neurons(sheet_x / um, sheet_y / um,
                                                             number_of_excitatory_neurons_per_row, tuning_map)

inhibitory_radial_axis = np.sqrt(inhibitory_axon_long_axis * inhibitory_axon_short_axis)

if optimization_label == "Repulsive":
    inhibitory_axonal_clouds = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_axon_long_axis / um,
                                                        inhibitory_axon_short_axis / um, sheet_x / um,
                                                        sheet_y / um, random_seed=2, n_iterations=1000)

    inhibitory_axonal_circles = create_interneuron_sheet_by_repulsive_force(N_I, inhibitory_radial_axis / um,
                                                         inhibitory_radial_axis / um, sheet_x / um,
                                                         sheet_y / um, random_seed=2, n_iterations=1000)
elif optimization_label == "Entropy Optimization":
    inhibitory_axonal_clouds, _ = create_interneuron_sheet_entropy_max_orientation(ex_positions, ex_tunings, N_I, inhibitory_axon_long_axis / um,
                                                        inhibitory_axon_short_axis / um, sheet_x / um,
                                                        sheet_y / um, trial_orientations=30)
    inhibitory_axonal_circles, _ = create_interneuron_sheet_entropy_max_orientation(ex_positions, ex_tunings, N_I, inhibitory_radial_axis / um,
                                                         inhibitory_radial_axis / um, sheet_x / um,
                                                         sheet_y / um, trial_orientations=1)


interneuron_tunings = [inhibitory_axonal_clouds, inhibitory_axonal_circles]
# plot_neural_sheet(ex_positions, ex_tunings, inhibitory_axonal_clouds)

'''
Neuron and synapse models
'''

excitatory_eqs = modellib.hodgkin_huxley_eqs_with_synaptic_conductance + modellib.eqs_ih
excitatory_params = modellib.hodgkin_huxley_params
excitatory_params.update(modellib.ih_params)
excitatory_params.update({"E_i": -80 * mV})
excitatory_params['ghbar'] = 0. * nS

lif_interneuron_eqs = """
dv/dt =1.0/tau* (-v + u_ext) :volt (unless refractory)
u_ext = u_ext_const : volt
"""

lif_interneuron_params = {
    "tau": 7 * ms,
    "v_threshold": -40 * mV,
    "v_reset": -60 * mV,
    "tau_refractory": 0.0 * ms,
    "u_ext_const": -50 * mV
}

lif_interneuron_options = {
    "threshold": "v>v_threshold",
    "reset": "v=v_reset",
    "refractory": "tau_refractory",
    'method': 'euler'
}

ei_synapse_model = modellib.delta_synapse_model
ei_synapse_on_pre = modellib.delta_synapse_on_pre
ei_synapse_param = modellib.delta_synapse_param

ie_synapse_model = modellib.exponential_synapse

ie_synapse_on_pre = modellib.exponential_synapse_on_pre
ie_synapse_param = modellib.exponential_synapse_params
ie_synapse_param["tau_syn"] = 2 * ms

'''
Connectvities
'''

# Spatial network with ellipsoid axons
ie_connections = get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_positions, inhibitory_axonal_clouds)

inhibitory_synapse_strength = 30 * nS
in_ex_weights = np.zeros((N_I, N_E)) * nS
for interneuron_idx, connected_excitatory_idxs in enumerate(ie_connections):
    in_ex_weights[interneuron_idx, connected_excitatory_idxs] = inhibitory_synapse_strength
#TODO: Remove this test
test_bins = range(10)
print(in_ex_weights.shape)
degree_of_exc_neuron = np.sum(in_ex_weights,axis=0)
print(degree_of_exc_neuron.shape)
plt.figure()
plt.hist(degree_of_exc_neuron/inhibitory_synapse_strength,test_bins)
plt.title('Ellipsoid exc degree')



excitatory_synapse_strength = 1 * mV
ex_in_weights = np.where(in_ex_weights > 0 * nS, excitatory_synapse_strength, 0 * mV).T * volt
print('ellipse ei_connections: ', np.sum(ex_in_weights.flatten())/excitatory_synapse_strength)
# Spatial network with circular axons
ie_connections_circle = get_excitatory_neurons_in_inhibitory_axonal_clouds(ex_positions, inhibitory_axonal_circles)

print('ellipse:',np.mean([len(row) for row in ie_connections]),' cirlce:',np.mean([len(row) for row in ie_connections_circle]))

in_ex_weights_circle = np.zeros((N_I, N_E)) * nS
for interneuron_idx, connected_excitatory_idxs in enumerate(ie_connections_circle):
    in_ex_weights_circle[interneuron_idx, connected_excitatory_idxs] = inhibitory_synapse_strength

#TODO: Remove this test

print(in_ex_weights_circle.shape)
degree_of_exc_neuron_circle = np.sum(in_ex_weights_circle,axis=0)
print(degree_of_exc_neuron_circle.shape)
plt.figure()
plt.hist(degree_of_exc_neuron_circle/inhibitory_synapse_strength,test_bins)
plt.title('Circle exc degree')
plt.show()



excitatory_synapse_strength = 1 * mV
ex_in_weights_circle = np.where(in_ex_weights_circle > 0 * nS, excitatory_synapse_strength, 0 * mV).T * volt
print('circle ei_connections: ', np.sum(ex_in_weights_circle.flatten())/excitatory_synapse_strength)

# No synapses
no_conn_ie = np.zeros((N_I, N_E)) * nS
no_conn_ei = np.zeros((N_E, N_I)) * mV

'''
Prepare nets
'''
nets = []

connectivity_label = ["No synapse", "Ellipsoid", "Circle"]
connectivities = [(no_conn_ei, no_conn_ie), (ex_in_weights, in_ex_weights), (ex_in_weights_circle, in_ex_weights_circle)]
for ei_weights, ie_weights in connectivities:
    net = ex_in_network(N_E, N_I, excitatory_eqs, excitatory_params, lif_interneuron_eqs,
                        lif_interneuron_params,
                        lif_interneuron_options, ei_synapse_model, ei_synapse_on_pre,
                        ei_synapse_param,
                        ei_weights, ie_synapse_model, ie_synapse_on_pre,
                        ie_synapse_param, ie_weights, random_seed=2)
    nets.append(net)

'''
Head direction input
'''

ex_input_baseline = 0.0 * nA

peak_phase = 0
input_sharpness = 1
direction_input = get_head_direction_input(peak_phase, input_sharpness)
max_head_direction_input_amplitude = 0.5 * nA
input_to_excitatory_population = ex_input_baseline + max_head_direction_input_amplitude * direction_input(
    np.array(ex_tunings))

'''
Run simulation
'''
skip = 100 * ms
length = 1100 * ms
duration = skip + length

for net in nets:
    excitatory_neurons = net["excitatory_neurons"]
    excitatory_neurons.I = input_to_excitatory_population
    net.run(duration)

'''
Get spatial map of rates
'''


def get_rates(spike_monitor, min_time=0 * ms):
    spike_trains = spike_monitor.spike_trains()

    isis = [np.ediff1d(np.extract(spike_times / ms > skip / ms, spike_times / ms)) * ms for spike_times in
            spike_trains.values()]

    rates = np.array([1.0 / np.mean(isi / ms) if isi.shape[0] != 0 else 0 for isi in isis]) * khertz
    return rates


excitatory_rates = [get_rates(net["excitatory_spike_monitor"]) for net in nets]

number_of_conditions = len(connectivities)

fig, axes = plt.subplots(1, number_of_conditions,figsize=(13.5,4.5))

X, Y = get_position_mesh(ex_positions)
for ax, net_label, rates in zip(axes, connectivity_label, excitatory_rates):
    c = ax.pcolor(X, Y, np.reshape(rates / hertz, (number_of_excitatory_neurons_per_row,
                                                   number_of_excitatory_neurons_per_row)), vmin=0, cmap='hot')
    ax.set_title(net_label)
    fig.colorbar(c, ax=ax, label="f (Hz)")

if save_dont_show:
    plt.savefig(output_folder + "/spatial_rate_map_{}.png".format(tuning_label))
else:
    plt.show()

'''
Get rate distribution over angles
'''
# Define model function to be used to fit to the data above:
def gauss(x, *p):
    A, mu, sigma, B = p
    return A*np.exp(-(x-mu)**2/(2.*sigma**2)) + B

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

from scipy.optimize import curve_fit



# Finally, lets get the fitting parameters, i.e. the mean and standard deviation:


rate_dist_bins = np.linspace(-np.pi,np.pi,40)

fig, ax = plt.subplots(1, 1)
coeff_list = []
# for rates, label in zip(excitatory_rates, connectivity_label):
#     tuning_mean, bin_edges, bin_number = binned_statistic(ex_tunings,rates / hertz,statistic='mean',bins=rate_dist_bins)
#     tuning_std_dev, _, _ = binned_statistic(ex_tunings, rates / hertz, statistic='std', bins=rate_dist_bins)
#     ax.errorbar(bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]), tuning_mean, tuning_std_dev, fmt='o', label=label)
#     # ax.scatter(ex_tunings, rates / hertz, label=label)
#     coeff, var_matrix = curve_fit(gauss, bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]), tuning_mean, p0=p0)
#     coeff_list.append(coeff)
#     print('Fitted mean = ', coeff[1])
#     print(label,'Fitted standard deviation = ', coeff[2])
# ax.plot(bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]),gauss(bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]),*coeff_list[1]),color='red',markersize='5')
# ax.plot(bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]),gauss(bin_edges[1:]-0.5*(bin_edges[1]-bin_edges[0]),*coeff_list[2]),color='blue',markersize='5')
ex_tunings_plt = np.array(ex_tunings)
sort_ids = ex_tunings_plt.argsort()
ex_tunings_plt = ex_tunings_plt[sort_ids]
for rates, label in zip(excitatory_rates, connectivity_label):
    rates_plt = np.array(rates)
    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)")

if save_dont_show:
    plt.savefig(output_folder + "/rate_over_angles_{}.png".format(tuning_label))
else:
    plt.show()

'''
Plot neural sheet for comparison
'''

for label, cloud in zip(connectivity_label[1:],interneuron_tunings):
    plot_neural_sheet(ex_positions, ex_tunings, cloud)
    if save_dont_show:
        plt.savefig(output_folder + "/neural_sheet_{}_{}.png".format(label,tuning_label))
    else:
        plt.show()

'''
Show spatio-temporal structure of activity
'''

dt_binning = 0.5 * ms
activity_decay_time_constant = 1.5 * ms

for label, net in zip(connectivity_label, nets):
    spike_trains = net["excitatory_spike_monitor"].spike_trains().values()

    times, activities = transform_spike_train_in_activity_array(spike_trains, dt_binning, duration,
                                                            activity_decay_time_constant)

    animation = create_neural_sheet_animation(times/ms, activities, frame_rate=25)
    animation.save(output_folder + "/animation_{}_{}.mp4".format(label,tuning_label))
    plt.show()