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=2, \
                                    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(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)

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

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

run_time = 200. * ms
exc_strength = lif_interneuron_params["v_threshold"]-lif_interneuron_params["v_reset"] + 5*mV

ex_drive_range = np.linspace(-0.1,0.1,10)*nA

input_current = 0.15*nA
inh_strength = 20.0 * nS

n_ghbar = 10
ghbar_range = linspace(0.0,50.0,n_ghbar)*nS
n_syn_strength = 10
synaptic_strength_range = linspace(0.0,160.0,n_syn_strength)*nS

gain_array = np.ndarray((n_ghbar,n_syn_strength), dtype=float)

syn_str_id = 0

for synaptic_strength_val in synaptic_strength_range:
    inh_strength = synaptic_strength_val
    ghbar_id = 0

    for ghbar_val in ghbar_range:
        rate_diff = []

        for drive in ex_drive_range:
            net.restore()
            excitatory_neurons[0].I = input_current + drive
            network_params['ghbar'] = ghbar_val
            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])
            rate_diff.append(rate_1-rate_2)

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

        print(ex_drive_range)
        print(rate_diff)
        gain = np.polyfit(ex_drive_range/nA,rate_diff/Hz,1)[0]
        print(gain)
        gain_array[ghbar_id, syn_str_id] = gain
        ghbar_id = ghbar_id + 1
    syn_str_id = syn_str_id + 1



'''
Plotting
'''

print(gain_array)
### Visualization of the h current role

x,y = np.meshgrid(ghbar_range / nS,synaptic_strength_range / nS)
# levels = MaxNLocator(nbins=2).tick_values(0.0,1.0)
# cmap = plt.get_cmap('PiYG')
# norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

fig, ax = plt.subplots(1, 1)

im = ax.pcolormesh(x,y,gain_array)#, cmap=cmap, norm=norm)
ax.set_title('bistability via iter. map')
ax.set_xlabel('ghbar (nS)')
ax.set_ylabel('syn_strength (nS)')
fig.colorbar(im, ax=ax)

fig.tight_layout()
plt.show()

