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




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

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)

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)

analyze_response_curve(period_unperturbed, inhibitory_delay, periods)

spike_recorder, neuron_state_recorder = timed_inhibition_experiment.run_sim(delay=3 * ms, record_states=True,
                                                                            pulse_strength=inhibition_off,
                                                                            duration=200 * ms)

### Visualization of the h current role


def plot_trace():
    t_state = neuron_state_recorder.t
    v_state = neuron_state_recorder[1].v
    g_state = neuron_state_recorder[1].g_syn
    r_state = neuron_state_recorder[1].r

    fig, (ax1, ax2, ax3) = plt.subplots(nrows=3)
    ax1.plot(t_state / ms, v_state / mV, 'C2')
    ax1.set_ylabel("V (mV)")

    ax2.plot(t_state / ms, r_state, 'C1')
    ax2.set_xlabel("t (ms)")
    ax2.set_ylabel("Ih activation")
    ax2.set_ylim(0,1)

    ax3.plot(t_state / ms, g_state/nS, 'C1')
    ax3.set_ylabel("g (nS)")

    fig.suptitle("Inhibition {:.1f}nS and delay {:.1f}ms".format(pulse_strength / nS, delay / ms))


delay = 0 * ms
pulse_strength = inhibition_off

spike_recorder, neuron_state_recorder = timed_inhibition_experiment.run_sim(delay=delay, record_states=True,
                                                                            pulse_strength=pulse_strength,
                                                                            duration=100 * ms)
plot_trace()

delay = 3 * ms
pulse_strength = inhibition_on

spike_recorder, neuron_state_recorder = timed_inhibition_experiment.run_sim(delay=delay, record_states=True,
                                                                            pulse_strength=pulse_strength,
                                                                            duration=100 * ms)
plot_trace()


delay = 15 * ms
pulse_strength = inhibition_on

spike_recorder, neuron_state_recorder = timed_inhibition_experiment.run_sim(delay=delay, record_states=True,
                                                                            pulse_strength=pulse_strength,
                                                                            duration=100 * ms)
plot_trace()

plt.show()