# Contrast the output of a ring network with unidirectional inhibition with one that has IE reciprocity

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

import scripts.models as modellib

from brian2 import seed

seed(2)

N_E = 100
ex_angles = np.arange(-np.pi, np.pi, 2 * np.pi / N_E)
N_I = 20
in_angles = np.arange(-np.pi, np.pi, 2 * np.pi / N_I)

'''
Excitatory neurons
'''

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_neurons = br.NeuronGroup(N_E, excitatory_eqs, namespace=excitatory_params,
                                    threshold='v>0*mV', refractory='v>0*mV', method='exponential_euler')

ex_input_baseline = 0.05 * nA
excitatory_neurons.I = ex_input_baseline
excitatory_neurons.v = modellib.hodgkin_huxley_params["El"] * np.random.normal(1, 0.1, N_E)

excitatory_spike_monitor = br.SpikeMonitor(excitatory_neurons)
recorded_excitatory_neuron_idx = int(N_E / 2)
excitatory_trace_recorder = br.StateMonitor(excitatory_neurons, 'v', record=[recorded_excitatory_neuron_idx])

excitatory_params['ghbar'] = 20. * nS

'''
Inhibitory neurons
'''

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

inhibitory_neurons = br.NeuronGroup(N_I, lif_interneuron_eqs, namespace=lif_interneuron_params,
                                    **lif_interneuron_options)

inhibitory_neurons.v = lif_interneuron_params["u_ext_const"]
recorded_interneuron_idx = int(N_I / 2)
inhibitory_trace_recorder = br.StateMonitor(inhibitory_neurons, 'v', record=[recorded_interneuron_idx])

inhibitory_spike_monitor = br.SpikeMonitor(inhibitory_neurons)

'''
Connectivities
'''


# EI Profile from Brunel: every excitatory neuron projects on inhibitory neurons of the same direction with a spread
# of 80 degrees

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

ex_in_differences = np.tile(ex_angles, (N_I, 1)).T - np.tile(in_angles, (N_E, 1))
ei_offset = 0
ei_offset_in_rad = ei_offset / 180.0 * np.pi
ei_width = 80.0 / 180.0 * np.pi
ei_weight_threshold = 0.1

max_excitatory_strength = 1 * mV

ex_in_weights = np.exp(-(get_phase_difference((ex_in_differences + ei_offset_in_rad))) ** 2 / 2 / ei_width ** 2)
ex_in_weights[np.where(ex_in_weights < ei_weight_threshold)] = 0.0

ex_indices, in_indices = ex_in_weights.nonzero()
ei_synapses = br.Synapses(source=excitatory_neurons, target=inhibitory_neurons, model=modellib.delta_synapse_model,
                          on_pre=modellib.delta_synapse_on_pre, namespace=modellib.delta_synapse_param)
ei_synapses.connect()

for ex_idx, in_idx in zip(ex_indices, in_indices):
    ei_synapses.synaptic_strength[ex_idx, in_idx] = max_excitatory_strength * ex_in_weights[ex_idx, in_idx]

ie_offset = 0
ie_offset_in_rad = ie_offset / 180.0 * np.pi
in_ex_difference = ex_in_differences.T
ie_width = 160.0 / 180.0 * np.pi
in_ex_weights = np.exp(
    -get_phase_difference(in_ex_difference + ie_offset_in_rad) ** 2 / 2 / ie_width ** 2)

ie_synapses = br.Synapses(source=inhibitory_neurons, target=excitatory_neurons, model=modellib.exponential_synapse,
                          on_pre=modellib.exponential_synapse_on_pre, namespace=modellib.exponential_synapse_params)

ie_synapses.connect()
ie_synapses.tau_syn = 2 * ms

max_inhibitory_strength = 50 * nS
in_indices, ex_indices = in_ex_weights.nonzero()
for in_idx, ex_idx in zip(in_indices, ex_indices):
    ie_synapses.synaptic_strength[in_idx, ex_idx] = max_inhibitory_strength * in_ex_weights[in_idx, ex_idx]

'''
Net
'''

net = br.Network(excitatory_neurons)
net.add(inhibitory_neurons)
net.add(ei_synapses)
net.add(ie_synapses)
net.add(excitatory_spike_monitor)
net.add(excitatory_trace_recorder)
net.add(inhibitory_spike_monitor)
net.add(inhibitory_trace_recorder)

'''
Transient head direction input
'''

start = 1000 * ms
length = 1000 * ms

phase_dispersion = 4
head_direction_input_amplitude = 1 * nA


def get_head_direction_input(peak_phase, phase_dispersion, excitatory_direction_preference):
    """
    Returns drive to a neuron with a certain preferred direction assuming that the input is distributed as a von Mises (
    Gaussian on a ring) distribution with a peak phase and an uncertainty given by the phase dispersion. Normalized
    to one at the peak phase.
    :param peak_phase:
    :param phase_dispersion:
    :param excitatory_direction_preference:
    :return:
    """
    head_direction_input_shape = np.exp(
        phase_dispersion * np.cos(excitatory_direction_preference - peak_phase)) / np.exp(phase_dispersion)
    return head_direction_input_shape


'''
Run simulation
'''

max_head_direction_input_amplitude = 0.5 * nA

first_peak_phase = np.pi/2.0
second_peak_phase = -np.pi / 2.0

input_dispersion = 3

length_of_input = 200 * ms
length_of_pause = 500 * ms


events = []


events.append((0*ms, 'Start'))
net.run(length_of_pause)


events.append((excitatory_trace_recorder.t[-1], "Input at {:.1f}".format(first_peak_phase/np.pi*180.0)))
excitatory_neurons.I = ex_input_baseline + max_head_direction_input_amplitude * get_head_direction_input(
    first_peak_phase,
    input_dispersion,
    ex_angles)
net.run(length_of_input)
events.append((excitatory_trace_recorder.t[-1], "No input".format(first_peak_phase/np.pi*180.0)))

excitatory_neurons.I = ex_input_baseline
net.run(length_of_pause)

events.append((excitatory_trace_recorder.t[-1], "Input at {:.1f}".format(second_peak_phase/np.pi*180.0)))
excitatory_neurons.I = ex_input_baseline + max_head_direction_input_amplitude * get_head_direction_input(
    second_peak_phase,
    input_dispersion,
    ex_angles)
net.run(length_of_input)

events.append((excitatory_trace_recorder.t[-1], "No input".format(first_peak_phase/np.pi*180.0)))
excitatory_neurons.I = ex_input_baseline
net.run(length_of_pause)

### Set synapses to zero to compare

events.append((excitatory_trace_recorder.t[-1], "Set synapses to zero".format(first_peak_phase/np.pi*180.0)))
ei_synapses.synaptic_strength = 0*mV
net.run(length_of_pause)

excitatory_neurons.I = ex_input_baseline + max_head_direction_input_amplitude * get_head_direction_input(
    first_peak_phase,
    input_dispersion,
    ex_angles)
net.run(length_of_input)

events.append((excitatory_trace_recorder.t[-1], "Input at {:.1f}".format(second_peak_phase/np.pi*180.0)))
excitatory_neurons.I = ex_input_baseline
net.run(length_of_pause)




'''
Plot
'''
duration = excitatory_trace_recorder.t[-1]

fig, axes = plt.subplots(5, 1, sharex=True)

ax_text = axes[0]
for t, comment in events:
    ax_text.text(t/ms, 0.5, " "+comment,fontsize=12 )
ax_text.axis('off')

ax_ex_trace = axes[1]
ax_ex_trace.plot(excitatory_trace_recorder.t / ms, excitatory_trace_recorder.v[0] / mV, 'r')
ax_ex_trace.set_ylabel("V ex (mV)")
ax_ex_trace.set_ylabel("E {:.0f}° (mV)".format(ex_angles[recorded_excitatory_neuron_idx]))

ax_in_trace = axes[2]
ax_in_trace.plot(inhibitory_trace_recorder.t / ms, inhibitory_trace_recorder.v[0] / mV, 'b')
ax_in_trace.set_ylabel("I {:.0f}° (mV)".format(in_angles[recorded_interneuron_idx]))

ax_ex = axes[3]
ax_ex.set_ylim(-180, 180)
ax_ex.set_ylabel("E angles")

ax_ex_raster = ax_ex.twinx()
ex_spike_trains = 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_raster.set_ylabel("Excitatory neurons")

ax_in = axes[4]
ax_in.set_ylim(-180, 180)
ax_in.set_ylabel("I angles")
ax_in_raster = ax_in.twinx()
in_spike_trains = 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|')

ax_in_raster.set_xlim(0, duration / ms)
ax_in_raster.set_ylim(0, N_I + 1)
ax_in_raster.set_ylabel("Inhibitory neurons")


for ax in axes:
    ymin, ymax =  ax.get_ylim()
    for t, _ in events:
        ax.vlines(t/ms, ymin, ymax, 'k', linestyles='--' )
plt.show()