import matplotlib.pyplot as plt
import numpy as np
from brian2.units import *

import scripts.models as modellib
from scripts.ring_network.head_direction import generate_ring_network, get_phase_difference, get_head_direction_input, get_half_width

'''
Neuron and synapse models
'''

N_E = 100
N_I = 20

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

### Non reciprocal

non_reciprocal_ei_offset = 0
non_reciprocal_ei_offset_in_rad = non_reciprocal_ei_offset / 180.0 * np.pi
non_reciprocal_ei_width = 30.0 / 180.0 * np.pi

non_reciprocal_excitatory_strength = 1.0 * mV


def non_reciprocal_ei(ex_in_differences):
    return (np.exp(-(get_phase_difference(
        (ex_in_differences + non_reciprocal_ei_offset_in_rad))) ** 2 / 2 / non_reciprocal_ei_width ** 2) + \
            np.exp(-(get_phase_difference(
                (ex_in_differences - non_reciprocal_ei_offset_in_rad))) ** 2 / 2 / non_reciprocal_ei_width **
                   2)) * non_reciprocal_excitatory_strength


non_reciprocal_ie_offset = 90
non_reciprocal_ie_offset_in_rad = non_reciprocal_ie_offset / 180.0 * np.pi
non_reciprocal_ie_width = 30.0 / 180.0 * np.pi

non_reciprocal_inhibitory_strength = 20 * nS


def non_reciprocal_ie(in_ex_difference):
    return (np.exp(-get_phase_difference(
        in_ex_difference + non_reciprocal_ie_offset_in_rad) ** 2 / 2 / non_reciprocal_ie_width ** 2) + \
            np.exp(-get_phase_difference(
                in_ex_difference - non_reciprocal_ie_offset_in_rad) ** 2 / 2 / non_reciprocal_ie_width **
                   2)) * non_reciprocal_inhibitory_strength


### Reciprocal

reciprocal_ei_offset = 45
reciprocal_ei_offset_in_rad = reciprocal_ei_offset / 180.0 * np.pi
reciprocal_ei_width = 20.0 / 180.0 * np.pi

reciprocal_excitatory_strength = 1.5 * mV


def reciprocal_ei(ex_in_differences):
    return (np.exp(-(get_phase_difference(
        (ex_in_differences + reciprocal_ei_offset_in_rad))) ** 2 / 2 / reciprocal_ei_width ** 2) + \
            np.exp(-(get_phase_difference(
                (ex_in_differences - reciprocal_ei_offset_in_rad))) ** 2 / 2 / reciprocal_ei_width **
                   2)) * reciprocal_excitatory_strength


reciprocal_ie_offset = 45
reciprocal_ie_offset_in_rad = reciprocal_ie_offset / 180.0 * np.pi
reciprocal_ie_width = 40.0 / 180.0 * np.pi

reciprocal_inhibitory_strength = 20 * nS


def reciprocal_ie(in_ex_difference):
    return (np.exp(-get_phase_difference(
        in_ex_difference + reciprocal_ie_offset_in_rad) ** 2 / 2 / reciprocal_ie_width ** 2) + \
            np.exp(-get_phase_difference(
                in_ex_difference - reciprocal_ie_offset_in_rad) ** 2 / 2 / reciprocal_ie_width **
                   2)) * reciprocal_inhibitory_strength


### No synapses

def no_conn(angular_difference):
    return np.zeros(np.shape(angular_difference))


### Plot connectivities


fig, axes = plt.subplots(1, 2, sharey=True)

angular_differences = np.arange(-np.pi, np.pi, 2 * np.pi / N_E)

ax = axes[0]
ax.set_title("Non-reciprocal")

ax.plot(angular_differences / np.pi * 180, non_reciprocal_ei(angular_differences) / np.max(non_reciprocal_ei(
    angular_differences)), 'r', label="E->I")
ax.plot(angular_differences / np.pi * 180,
        non_reciprocal_ie(angular_differences) / np.max(non_reciprocal_ie(angular_differences)), 'b',
        label="I->E")

ax = axes[1]
ax.set_title("Reciprocal")
ax.plot(angular_differences / np.pi * 180, reciprocal_ei(angular_differences) / np.max(reciprocal_ei(
    angular_differences)), 'r', label="E->I")
ax.plot(angular_differences / np.pi * 180,
        reciprocal_ie(angular_differences) / np.max(reciprocal_ie(angular_differences)), 'b',
        label="I->E")

for ax in axes:
    ax.set_xlabel("Angle")
    ax.legend()
axes[0].set_ylabel("w/w$_{max}$")

'''
Prepare nets
'''

connectivity_label = ["No synapse", "Non-reciprocal", "Reciprocal"]
connectivities = [(no_conn, no_conn), (non_reciprocal_ei, non_reciprocal_ie), (reciprocal_ei, reciprocal_ie)]

nets = []

for ei, ie in connectivities:
    ex_angles, in_angles, net = generate_ring_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, ie_synapse_model, ie_synapse_on_pre,
                                                      ie_synapse_param, ie, random_seed=2)
    nets.append(net)

'''
Head direction input
'''

ex_input_baseline = 0.0 * nA

first_peak_phase = 0
input_sharpness = 1
direction_input = get_head_direction_input(first_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(ex_angles)

half_width_input = get_half_width(ex_angles, input_to_excitatory_population)

'''
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 tuning of the network
'''

fig, ax = plt.subplots(1, 1)
output_handles = []
output_widths = []

for net_label, net in zip(connectivity_label, nets):
    excitatory_spike_monitor = net["excitatory_spike_monitor"]
    spike_trains = excitatory_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
    half_width_output = get_half_width(ex_angles, rates)

    output_widths.append(half_width_output / np.pi * 180)
    output_handles.append(ax.plot(ex_angles / np.pi * 180, rates / hertz))

ax.set_xlabel("Angle")
ax.set_ylabel("f(Hz)")

ax_input = ax.twinx()
input_label = "Input width {" \
              ":.1f}°".format(half_width_input / np.pi * 180)
input_handle = ax_input.plot(ex_angles / np.pi * 180, input_to_excitatory_population / nA, 'k--')

handles = [input_handle[0]] + [output_handle[0] for output_handle in output_handles]
labels = [input_label] + ["{:s} output width {:.1f}°".format(type, width) for type, width in zip(connectivity_label,
                                                                                                output_widths)]
ax.legend(handles, labels)
ax_input.set_ylabel("Input (nA)")

'''
Spike trains
'''

fig, axes = plt.subplots(8, 1, sharex=True, gridspec_kw={"height_ratios": [1, 1, 0.1, 1, 1, 0.1, 1, 1]})

axs = [(axes[0], axes[1]), (axes[3], axes[4]), (axes[6], axes[7])]

for net_label, net, raster_axes in zip(connectivity_label, nets, axs):

    ax_ex = raster_axes[0]
    ax_in = raster_axes[1]

    ax_ex.set_ylim(-180, 180)
    ax_ex.set_ylabel("E angles")

    ax_ex_raster = ax_ex.twinx()
    ex_spike_trains = net["excitatory_spike_monitor"].spike_trains()
    for neuron_idx, spike_times in ex_spike_trains.items():
        ax_ex_raster.plot(spike_times / ms, neuron_idx + 1 * np.ones(spike_times.shape), 'r|')

    ax_ex_raster.set_xlim(0, duration / ms)
    ax_ex_raster.set_ylim(0, N_E + 1)
    ax_ex.set_title(net_label)

    ax_in.set_ylim(-180, 180)
    ax_in.set_ylabel("I angles")
    ax_in_raster = ax_in.twinx()
    in_spike_trains = net["inhibitory_spike_monitor"].spike_trains()
    for neuron_idx, spike_times in in_spike_trains.items():
        ax_in_raster.plot(spike_times / ms, neuron_idx + 1 * np.ones(spike_times.shape), 'b|')

axes[-1].set_xlabel("Time (ms)")

for ax in (axes[2], axes[5]):
    ax.set_frame_on(False)
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)

plt.show()
