import copy
import warnings

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

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', 'ih', '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])

'''
External Stimulus
'''
#First Try

# @implementation('cython', '''
#     cdef double stimulus_fun(double t):
#         return 6.25*1e-9 * sin(4. * 3.14 /(1000. 1e-3) * t)
#     ''')
# @check_units(t=second, result=amp)
# def stimulus_fun(t):
#     # return 4 * nA *t /ms
#     return 6.25* nA * sin(4. * 3.14 /(1000 * ms) * t)

#Second Try

# pulse = 0.25
# stimulus_array = br.TimedArray([0.,0.,0.,0.,0.,0.,pulse,0.,0.,0.,0.,0.,0.,-pulse,0.,0.,0.,0.,0.,0.] * nA,dt=100.*ms)

#Third Try
pulse_length = 25*ms
pulse_1_start = 200*ms
pulse_2_start = 400*ms
pulse_strength = 0.6 * nA
@network_operation(dt=1*ms)
def update_input(t):
    if t > pulse_1_start and t < pulse_1_start + 2*ms:
        excitatory_neurons[0].I = pulse_strength
    elif t > pulse_1_start + pulse_length and t < pulse_1_start + pulse_length + 2*ms:
        excitatory_neurons[0].I = input_current
    elif t > pulse_2_start and t < pulse_2_start + 2*ms:
        excitatory_neurons[1].I = pulse_strength
    elif t > pulse_2_start + pulse_length and t < pulse_2_start + pulse_length + 2*ms:
        excitatory_neurons[1].I = input_current



net = br.Network(update_input)

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



inh_strength = 150 * nS
exc_strength = lif_interneuron_params["v_threshold"]-lif_interneuron_params["v_reset"] + 5*mV

run_sim(inh_strength=inh_strength, exc_strength=exc_strength, record_states=True, run_time=run_time)
plot_spiking()
print(neuron_state_recorder.v.shape)
t_state = neuron_state_recorder.t
v1_state = neuron_state_recorder[0].v
v2_state = neuron_state_recorder[1].v
I_state_1 = neuron_state_recorder[0].I
I_state_2 = neuron_state_recorder[1].I

fig, axs = plt.subplots(3, sharex=True)
axs[0].plot(t_state / ms, v1_state / mV, 'C2')
axs[1].plot(t_state / ms, v2_state / mV, 'C1')
axs[2].plot(t_state / ms, I_state_1 / nA, 'C2')
axs[2].plot(t_state / ms, I_state_2 / nA, 'C1')
for ax in axs[:2]:
    ax.set_ylim(-100, 50)

run_sim(inh_strength=0*nS, exc_strength=10.*mV, record_states=True, run_time=run_time)
plot_spiking()

plt.show()





plt.show()

