import matplotlib.pyplot as plt
from brian2.units import *

from scripts.models import hodgkin_huxley_params, hodgkin_huxley_eqs_with_synaptic_conductance, ih_params, eqs_ih, \
    exponential_synapse, exponential_synapse_on_pre, exponential_synapse_params
from scripts.prc_and_iterative_diagram.timed_inhibition import TimedInhibition, bistability_from_iterative_map_merged
from numpy import ndarray,meshgrid

### 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": -120 * mV})

spike_threshold = -40 * mV

inhibition_off = 0.0 * nS
inhibition_on = 150 * nS

record_variables = ['v', 'r', 'g_syn']
integration_method = 'exponential_euler'

initial_states = {
    "v": hodgkin_huxley_params["El"]
}
current_drive = 0.15 * nA

n_simulations = 50
before_spike = 1 * ms
after_spike = 1 * ms

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

bistable_array = ndarray((n_ghbar,n_syn_strength), dtype=float)

# Incredibly, enumerate doesn't work with Quantity objects
syn_str_id = 0

timed_inhibition_experiment = TimedInhibition(neuron_eqs, synapse_model, synapse_on_pre, neuron_params,
                                                      synapse_params, initial_states, spike_threshold, integration_method,
                                                      current_drive, record_variables)

for synaptic_strength_val in synaptic_strength_range:
    inhibition_on = synaptic_strength_val
    ghbar_id = 0
    for ghbar_val in ghbar_range:
        timed_inhibition_experiment.net.restore()
        timed_inhibition_experiment.network_params['ghbar'] = ghbar_val
        timed_inhibition_experiment.net.store()


        period_unperturbed, inhibitory_delay, periods = timed_inhibition_experiment.measure_response_curve(n_simulations,
                                                                                                           inhibition_on,
                                                                                                           offset_before_spike_peak=before_spike,
                                                                                                           offset_after_spike_peak=after_spike)

        isbistable = bistability_from_iterative_map_merged(period_unperturbed, inhibitory_delay, periods)
        print(synaptic_strength_val,ghbar_val,isbistable)

        bistable_array[ghbar_id,syn_str_id] = isbistable
        ghbar_id = ghbar_id + 1
    syn_str_id = syn_str_id + 1

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

x,y = 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,bistable_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()