import numpy as np
import brian2 as br
from brian2 import mV, ms
import matplotlib.pyplot as plt
from scripts.spatial_network.head_direction_index_over_noise_scale import lif_interneuron_eqs, \
    lif_interneuron_params, lif_interneuron_options, calculate_rates

def inh_spike_rate_by_input(input_to_inhibitory_neuron):
    inhibitory_neurons = br.NeuronGroup(1, lif_interneuron_eqs, namespace=lif_interneuron_params, name="interneurons",
                                        **lif_interneuron_options)
    inhibitory_neurons.v = lif_interneuron_params["u_ext_const"]
    recorded_interneuron_idx = 0
    inhibitory_trace_recorder = br.StateMonitor(inhibitory_neurons, 'v', record=[recorded_interneuron_idx],
                                                name="inhibitory_trace_recorder")
    inhibitory_spike_monitor = br.SpikeMonitor(inhibitory_neurons, name="inhibitory_spike_monitor")


    # inhibitory_synapse_strength = traj.synapse.inhibitory * nS
    # excitatory_synapse_strength = traj.synapse.excitatory * mV

    net = br.Network(inhibitory_neurons)
    net.add(inhibitory_spike_monitor)
    net.add(inhibitory_trace_recorder)
    net.store()

    inhibitory_neurons = net["interneurons"]
    inhibitory_neurons.u_ext = input_to_inhibitory_neuron * mV
    # inhibitory_neurons.u_ext = traj.inh_input.baseline * mV

    inhibitory_neurons.tau = 30. * ms

    net.run(1000 * ms)
    spike_times = inhibitory_spike_monitor.spike_trains()
    print(spike_times)
    spike_rate = calculate_rates(spike_times.values())
    return spike_rate



def main():

    input_range = np.linspace(-42.0, 40.0, 21, endpoint=True)
    spike_rates = []
    for input in input_range:
        spike_rate = inh_spike_rate_by_input(input)
        spike_rates.append(spike_rate)

    plt.plot(input_range, spike_rates)
    plt.show()

    # traj.f_add_parameter_group("synapse")
    # traj.f_add_parameter("synapse.inhibitory", 1.0, "Strength of conductance-based inhibitory synapse in nS.")
    # traj.f_add_parameter("synapse.excitatory", 2.5, "Strength of conductance-based inhibitory synapse in mV.")
    #
    # traj.f_add_parameter("input.baseline", 0.05, comment="Head direction input baseline")
    # traj.f_add_parameter("input.amplitude", 0.6, comment="Head direction input amplitude")
    #
    # traj.f_add_parameter_group("inh_input")
    # traj.f_add_parameter("inh_input.baseline", -38., comment="Head direction input baseline")
    # traj.f_add_parameter("inh_input.amplitude", 18., comment="Head direction input amplitude")
    #
    # traj.f_add_parameter_group("interneuron")
    # traj.f_add_parameter("interneuron.tau", 30., comment="Interneuron timescale in ms")
    #
    #
    # traj.f_add_parameter_grou<p("simulation")
    # traj.f_add_parameter("simulation.entropy_maximisation.steps", 30, comment="Steps for entropy maximisation")
    # traj.f_add_parameter("simulation.dt", 0.1, comment="Network simulation time step in ms")
    # traj.f_add_parameter("simul>ation.duration", 1000, comment="Network simulation duration in ms")
    #
    #     "synapse.inhibitory": [5.],
    #     "synapse.excitatory": [0.001]


if __name__ == "__main__":
    main()