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


network_params = copy.deepcopy(excitatory_neuron_params)
network_params.update(lif_interneuron_params)
network_params["spike_threshold"] = spike_threshold
record_variables = ['v', 'm', 'n', 'h']
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
inhibitory_neurons.v = lif_interneuron_params["v_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()

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

    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)

'''
Extract LC and map phase to it
'''

run_time = 60. * ms
inh_strength = 0 * 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)

t_state = neuron_state_recorder.t
v1_state = neuron_state_recorder[0].v
m1_state = neuron_state_recorder[0].m
n1_state = neuron_state_recorder[0].n
h1_state = neuron_state_recorder[0].h

dt = t_state[-1] / len(t_state)
t_0 = 25. * ms
id_1 = int(t_0 / dt)
print(t_state[id_1])
for i in range(id_1+1, len(t_state)):
    if v1_state[i]-v1_state[id_1] < 0.01 * v1_state[i] and m1_state[i]-m1_state[id_1] < 0.01 * m1_state[i]:
        id_end = i
print(t_state[id_end])

phase_1 = 0.09
id_2 = id_1 + int(phase_1 * (id_end - id_1))

initial_state_1 = {
    "v": neuron_state_recorder[0].v[id_1],
    'm': neuron_state_recorder[0].m[id_1],
    'n': neuron_state_recorder[0].n[id_1],
    'h': neuron_state_recorder[0].h[id_1]
}

initial_state_2 = {
    "v": neuron_state_recorder[0].v[id_2],
    'm': neuron_state_recorder[0].m[id_2],
    'n': neuron_state_recorder[0].n[id_2],
    'h': neuron_state_recorder[0].h[id_2]
}
print(initial_state_1)
print(initial_state_2)
net.restore()
set_parameters_from_dict(excitatory_neurons[0],initial_state_1)
set_parameters_from_dict(excitatory_neurons[1],initial_state_2)

run_time = 600. * 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()

t_state = neuron_state_recorder.t
v1_state = neuron_state_recorder[0].v
v2_state = neuron_state_recorder[1].v
m1_state = neuron_state_recorder[0].m
m2_state = neuron_state_recorder[1].m
n1_state = neuron_state_recorder[0].n
n2_state = neuron_state_recorder[1].n
h1_state = neuron_state_recorder[0].h
h2_state = neuron_state_recorder[1].h

fig, axs = plt.subplots(4, sharex=True)
axs[0].plot(t_state / ms, v1_state / mV, 'C2')
axs[0].plot(t_state / ms, v2_state / mV, 'C1')
axs[1].plot(t_state / ms, m1_state, 'C2')
axs[1].plot(t_state / ms, m2_state, 'C1')
axs[2].plot(t_state / ms, n1_state, 'C2')
axs[2].plot(t_state / ms, n2_state, 'C1')
axs[3].plot(t_state / ms, h1_state, 'C2')
axs[3].plot(t_state / ms, h2_state, 'C1')


plt.show()

