import copy
import warnings

import brian2 as br
import matplotlib.pyplot as plt
import numpy as np
from brian2.units import *
from scripts.prc_and_iterative_diagram.timed_inhibition import get_mean_period

from scripts.models import hodgkin_huxley_eqs_with_synaptic_conductance, exponential_synapse, \
    exponential_synapse_on_pre, exponential_synapse_params, eqs_ih, hodgkin_huxley_params, \
    ih_params, delta_synapse_model, delta_synapse_on_pre


def set_parameters_from_dict(neurongroup, dictionary_of_parameters):
    for param_key, param_value in dictionary_of_parameters.items():
        try:
            neurongroup.__setattr__(param_key, param_value)
        except AttributeError as err:
            warnings.warn("{:s} has no parameter {:s}".format(neurongroup.name, param_key))

'''
Model for the excitatory neurons
'''
# Hodgkin Huxley model from Brian2 documentation

spike_threshold = 40*mV

excitatory_neuron_options = {
    "threshold" : "v > spike_threshold",
    "refractory" : "v > spike_threshold",
    "method" : 'exponential_euler'
}


'''
Model for the interneuron, currently only the inhibition timing is of interest thus the lif model
'''

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": -50*mV,
    "tau_refractory": 0.0*ms,
    "u_ext_const" : - 41 * mV
}

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



'''
Synapse Models
'''

delta_synapse_delay = 2.0 * ms


'''
Setup of the Model
'''

## With voltage delta synapse


ei_synapse_options = {
    'model' : delta_synapse_model,
    'on_pre' : delta_synapse_on_pre,
}



#
# ie_synapse_options = {
#     'model' : delta_synapse_model,
#     'on_pre' : delta_synapse_on_pre,
#     'delay' : delta_synapse_delay
# }


excitatory_neuron_eqs = hodgkin_huxley_eqs_with_synaptic_conductance + eqs_ih
excitatory_neuron_params = hodgkin_huxley_params
excitatory_neuron_params.update(ih_params)
excitatory_neuron_params.update(excitatory_neuron_params)
excitatory_neuron_params.update({"E_i": -120 * mV})

# excitatory_neuron_params["stimulus"] = stimulus_array
# excitatory_neuron_params["stimulus"] = stimulus_fun


# ### With conductance based delta synapse
# neuron_eqs = hodgkin_huxley_eqs_with_synaptic_conductance + eqs_ih
#
# synapse_model = exponential_synapse
# synapse_on_pre = exponential_synapse_on_pre
# synapse_params = exponential_synapse_params
#
# neuron_params = hodgkin_huxley_params
# neuron_params.update(ih_params)
# neuron_params.update({"E_i": -80 * mV})
#
# inhibition_off = 0.0 * nS
# inhibition_on = 100 * nS

network_params = copy.deepcopy(excitatory_neuron_params)
network_params.update(lif_interneuron_params)
network_params["spike_threshold"] = spike_threshold
record_variables = ['v', 'I']
integration_method = 'exponential_euler'

initial_states = {
    "v": hodgkin_huxley_params["El"]
}

threshold_eqs = """
spike_threshold: volt
"""

'''
Setup of the neuron groups, synapses and ultimately the network 
'''

excitatory_neurons = br.NeuronGroup(N=2, \
                                    model=excitatory_neuron_eqs, \
                                    threshold=excitatory_neuron_options["threshold"], \
                                    refractory=excitatory_neuron_options["refractory"], \
                                    method=excitatory_neuron_options["method"])
# set_parameters_from_dict(excitatory_neurons, excitatory_neuron_params) Why doesnt this work?
set_parameters_from_dict(excitatory_neurons, initial_states)
# excitatory_neurons.I_ext_const = 0.15*nA
input_current = 0.15*nA
excitatory_neurons.I = input_current




inhibitory_neurons = br.NeuronGroup(N=1, \
                                    model=lif_interneuron_eqs, \
                                    threshold=lif_interneuron_options["threshold"], \
                                    refractory=lif_interneuron_options["refractory"], \
                                    reset=lif_interneuron_options["reset"])
# set_parameters_from_dict(inhibitory_neurons, lif_interneuron_params) Same here

e_to_i_synapses = br.Synapses(source=excitatory_neurons, \
                              target=inhibitory_neurons, \
                              model=ei_synapse_options["model"], \
                              on_pre=ei_synapse_options["on_pre"])
e_to_i_synapses.connect()

i_to_e_synapses = br.Synapses(source=inhibitory_neurons, \
                              target=excitatory_neurons, \
                              model=exponential_synapse, \
                              on_pre=exponential_synapse_on_pre, \
                              delay=2.0*ms)
i_to_e_synapses.connect()
set_parameters_from_dict(i_to_e_synapses,exponential_synapse_params)

exc_spike_recorder = br.SpikeMonitor(source=excitatory_neurons)
inh_spike_recorder = br.SpikeMonitor(source=inhibitory_neurons)
neuron_state_recorder = br.StateMonitor(excitatory_neurons, record_variables, record=[0,1])


net = br.Network()

net.add(excitatory_neurons)
net.add(inhibitory_neurons)
net.add(e_to_i_synapses)
net.add(i_to_e_synapses)
net.add(exc_spike_recorder)
net.add(inh_spike_recorder)
net.add(neuron_state_recorder)
net.store()


'''
Run of the simulation
'''

def nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return array[idx]

def plot_spiking():
    ex_spike_trains = exc_spike_recorder.spike_trains()
    in_spike_trains = inh_spike_recorder.spike_trains()

    fig = plt.figure()
    ax = fig.add_subplot(111)
    for key, times in ex_spike_trains.items():
        ax.plot(times / ms, key / 2.0 * np.ones(times.shape), 'b|')

    offset = 2
    for key, times in in_spike_trains.items():
        ax.plot(times / ms, (key + offset) / 2.0 * np.ones(times.shape), 'r|')

    ax.grid(axis='x')
    ax.set_ylim(-0.1, 1.1)
    ax.set_xlabel("Time(ms)");


def run_sim(inh_strength, exc_strength, record_states=True, run_time=50 * ms):
    net.restore()

    e_to_i_synapses.synaptic_strength = exc_strength
    i_to_e_synapses.synaptic_strength = inh_strength

    # excitatory_neurons.stimulus = get_sinoid_stimulus(0.25 * nA, 4 * np.pi / run_time)
    # network_params["omega"] = 4 * np.pi / run_time
    # excitatory_neurons.I_ext_diff = [0.25, -0.25] * nA

    # neuron_state_recorder.record = record_states

    net.run(duration=run_time, namespace=network_params)

run_time = 1000. * ms
network_params['ghbar'] = 0.0 * nS
exc_strength = lif_interneuron_params["v_threshold"]-lif_interneuron_params["v_reset"] + 5*mV

ex_drive_range = np.linspace(-0.1,0.3,200)
# drive_for_traces = np.array([0.0,0.5,0.75])
drive_for_traces = np.array([])
ex_drive_range = np.append(ex_drive_range,drive_for_traces)*nA
drive_for_traces = drive_for_traces * nA

inhibition_on = 20.0 * nS

t_saved = False
t_values = []
v_values = []

'''
Gain in relation to a disconnected neuron
'''


input_current = 0.15*nA
inh_strength = 0.0 * nS
rate_diff_disconnected = []
rates_disconnected = []
for drive in ex_drive_range:
    net.restore()
    excitatory_neurons[0].I = input_current + drive
    net.store()

    run_sim(inh_strength=inh_strength, exc_strength=exc_strength, record_states=True, run_time=run_time)

    ex_spike_trains = exc_spike_recorder.spike_trains()
    if len(ex_spike_trains[0]) <= 1:
        rate_1 = 0.0
    else:
        rate_1 = 1.0/get_mean_period(ex_spike_trains[0])
    if len(ex_spike_trains[1]) <= 1:
        rate_2 = 0.0
    else:
        rate_2 = 1.0 / get_mean_period(ex_spike_trains[1])
    rates_disconnected.append([rate_1,rate_2])
    rate_diff_disconnected.append(rate_1-rate_2)

    if drive/nA == nearest(ex_drive_range/nA,0.0):
        print('drive: ',drive,"rates: ",rate_1,rate_2)

    if not t_saved:
        t_values = neuron_state_recorder.t
        t_saved = True

    if drive in drive_for_traces:
        print("drive for trace: ",drive)
        v_values.append(neuron_state_recorder.v[0])
        v_values.append(neuron_state_recorder.v[1])



'''
Gain in merged microcircuit
'''

input_current = 0.15*nA
inh_strength = inhibition_on
rate_diff_merged = []
rates_merged = []

for drive in ex_drive_range:
    net.restore()
    excitatory_neurons[0].I = input_current + drive
    net.store()

    run_sim(inh_strength=inh_strength, exc_strength=exc_strength, record_states=True, run_time=run_time)

    ex_spike_trains = exc_spike_recorder.spike_trains()
    if len(ex_spike_trains[0]) <= 1:
        rate_1 = 0.0
    else:
        rate_1 = 1.0/get_mean_period(ex_spike_trains[0])
    if len(ex_spike_trains[1]) <= 1:
        rate_2 = 0.0
    else:
        rate_2 = 1.0 / get_mean_period(ex_spike_trains[1])
    rates_merged.append([rate_1, rate_2])
    rate_diff_merged.append(rate_1-rate_2)

    if drive/nA == nearest(ex_drive_range/nA,0.0):
        print("rates: ",rate_1,rate_2)

    if drive in drive_for_traces:
        print("drive for trace: ",drive)
        v_values.append(neuron_state_recorder.v[0])
        v_values.append(neuron_state_recorder.v[1])

'''
Gain in unidirectional microcircuit
'''
net.restore()
net.remove(inhibitory_neurons)
net.remove(i_to_e_synapses)
net.remove(e_to_i_synapses)
net.remove(inh_spike_recorder)

inhibitory_neurons = br.NeuronGroup(N=2, \
                                    model=lif_interneuron_eqs, \
                                    threshold=lif_interneuron_options["threshold"], \
                                    refractory=lif_interneuron_options["refractory"], \
                                    reset=lif_interneuron_options["reset"])

e_to_i_synapses = br.Synapses(source=excitatory_neurons, \
                              target=inhibitory_neurons, \
                              model=ei_synapse_options["model"], \
                              on_pre=ei_synapse_options["on_pre"])
e_to_i_synapses.connect(condition='i == j')

i_to_e_synapses = br.Synapses(source=inhibitory_neurons, \
                              target=excitatory_neurons, \
                              model=exponential_synapse, \
                              on_pre=exponential_synapse_on_pre, \
                              delay=2.0*ms)
i_to_e_synapses.connect(condition='i != j')
set_parameters_from_dict(i_to_e_synapses,exponential_synapse_params)

inh_spike_recorder = br.SpikeMonitor(source=inhibitory_neurons)

net.add(inhibitory_neurons)
net.add(i_to_e_synapses)
net.add(e_to_i_synapses)
net.add(inh_spike_recorder)
net.store()
print(e_to_i_synapses.i,e_to_i_synapses.j)
print(i_to_e_synapses.i,i_to_e_synapses.j)


input_current = 0.15*nA
inh_strength = inhibition_on
rate_diff_unidirectional = []
rates_unidirectional = []

for drive in ex_drive_range:
    net.restore()
    excitatory_neurons[0].I = input_current + drive
    net.store()

    run_sim(inh_strength=inh_strength, exc_strength=exc_strength, record_states=True, run_time=run_time)

    ex_spike_trains = exc_spike_recorder.spike_trains()
    if len(ex_spike_trains[0]) <= 1:
        rate_1 = 0.0
    else:
        rate_1 = 1.0/get_mean_period(ex_spike_trains[0])
    if len(ex_spike_trains[1]) <= 1:
        rate_2 = 0.0
    else:
        rate_2 = 1.0 / get_mean_period(ex_spike_trains[1])
    rates_unidirectional.append([rate_1, rate_2])
    rate_diff_unidirectional.append(rate_1-rate_2)

    if drive/nA == nearest(ex_drive_range/nA,0.0):
        print("rates: ",rate_1,rate_2)

    if drive in drive_for_traces:
        print("drive for trace: ",drive)
        v_values.append(neuron_state_recorder.v[0])
        v_values.append(neuron_state_recorder.v[1])

rates_disconnected = np.array(rates_disconnected)
rates_merged = np.array(rates_merged)
rates_unidirectional = np.array(rates_unidirectional)

fig, axs = plt.subplots(2, sharex=True, sharey=True)
axs[0].plot(ex_drive_range / nA, rate_diff_disconnected / Hz, 'C1',label='disconnected')
axs[0].plot(ex_drive_range / nA, rate_diff_merged / Hz, 'C2', label='merged')
axs[0].plot(ex_drive_range / nA, rate_diff_unidirectional / Hz, 'C3', label='unidirectional')

axs[1].plot(ex_drive_range / nA,rates_disconnected[:,0] / Hz,'C1')
axs[1].plot(ex_drive_range / nA,rates_disconnected[:,1] / Hz,'C1',linestyle='--')

axs[1].plot(ex_drive_range / nA,rates_merged[:,0] / Hz,'C2')
axs[1].plot(ex_drive_range / nA,rates_merged[:,1] / Hz,'C2',linestyle='--')

axs[1].plot(ex_drive_range / nA,rates_unidirectional[:,0] / Hz,'C3')
axs[1].plot(ex_drive_range / nA,rates_unidirectional[:,1] / Hz,'C3',linestyle='--')

axs[1].set_xlabel('Input Current [nA]')
axs[0].set_ylabel('Firing Rate [Hz]')
axs[1].set_ylabel('Firing Rate [Hz]')
axs[0].legend()
# run_sim(inh_strength=0*nS, exc_strength=10.*mV, record_states=True, run_time=run_time)
# plot_spiking()
#
# plt.show()
v_values = np.array(v_values)
print(v_values.shape)


fig, axs = plt.subplots(3*len(drive_for_traces), sharex=True, sharey=True)
for i,ax in enumerate(axs):
    ax.plot(t_values / ms, v_values[int(i*2)] / mV, 'C1', label='drive: ' + str(drive_for_traces[i%len(drive_for_traces)]/nA) + " nA")
    ax.plot(t_values / ms, v_values[int(i*2+1)] / mV, 'C2',label='passive')
    ax.legend(loc='right')



plt.show()

