import copy
import warnings

import brian2 as br
import matplotlib.pyplot as plt
import numpy as np
from brian2.units import *


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


def get_mean_period(spike_train):
    return (np.max(spike_train) - np.min(spike_train)) / (spike_train.shape[0] - 1)


def get_difference_in_periods(spike_trains_dict, t_isi):
    mean_response_dict = {}
    #     t_isi = get_mean_period(spike_trains_dict[0.0 * ms])
    for delay, spike_train in spike_trains_dict.items():
        mean_response_dict[delay] = get_mean_period(spike_train) - t_isi
    return mean_response_dict


def plot_spiking(spike_trains_dict):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    for key, times in spike_trains_dict.items():
        ax.plot(times / ms, key * np.ones(times.shape), 'b.')
    #     ax.plot(spike_train[0]/ms, np.ones(spike_train[0].shape), 'b|')

    ax.grid(axis='x')
    #     ax.set_ylim(-0.1, 1.1)
    ax.set_xlabel("Time(ms)")


def plot_mean_period(mean_period_dict):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(list(mean_period_dict.keys()), list(mean_period_dict.values()), 'b')

    ax.grid(axis='x')
    ax.set_xlabel("Delay (ms)")


def plot_mean_response(mean_response_dict):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(np.array(list(mean_response_dict.keys())), 1e3 * np.array(list(mean_response_dict.values())), 'b')
    ax.plot(list(mean_response_dict.keys()), list(mean_response_dict.keys()), 'k--')

    ax.grid(axis='x')
    ax.set_xlabel("t_inh (ms)")
    ax.set_ylabel("exc_spike_delay (ms)")


def plot_delay_minus_response(mean_response_dict):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    print(np.array(list(mean_response_dict.keys())))
    print(np.array(list(mean_response_dict.values())))
    print(np.array(list(mean_response_dict.keys())) - np.array(list(mean_response_dict.values())))
    ax.plot(list(mean_response_dict.keys()),
            np.array(list(mean_response_dict.keys())) - 1e3 * np.array(list(mean_response_dict.values())), 'b')

    ax.grid(axis='x')
    ax.set_xlabel("Delay (ms)");


def plot_phase_response_curve(mean_response_dict, t_isi):
    phi_inh = np.array(list(mean_response_dict.keys())) * ms / t_isi
    delta_phi = -np.array(list(mean_response_dict.values())) * second / t_isi

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(phi_inh, delta_phi, 'b')
    ax.plot(phi_inh, -phi_inh, 'k--')

    ax.grid(axis='x')
    ax.set_xlabel("phase of inhibition (ms)")
    ax.set_ylabel("phase shift (ms)");


# Hodgkin Huxley model from Brian2 documentation
area = 20000 * umetre ** 2

hodgkin_huxley_params = {
    "Cm": 1 * ufarad * cm ** -2 * area,
    "gl": 5e-5 * siemens * cm ** -2 * area,
    "El": -65 * mV,
    "EK": -90 * mV,
    "ENa": 50 * mV,
    "g_na": 100 * msiemens * cm ** -2 * area,
    "g_kd": 30 * msiemens * cm ** -2 * area,
    "VT": -63 * mV
}
# The model
hodgkin_huxley_eqs = br.Equations('''
dv/dt = (gl*(El-v) - g_na*(m*m*m)*h*(v-ENa) - g_kd*(n*n*n*n)*(v-EK) + I + ih)/Cm : volt
dm/dt = 0.32*(mV**-1)*(13.*mV-v+VT)/
    (exp((13.*mV-v+VT)/(4.*mV))-1.)/ms*(1-m)-0.28*(mV**-1)*(v-VT-40.*mV)/
    (exp((v-VT-40.*mV)/(5.*mV))-1.)/ms*m : 1
dn/dt = 0.032*(mV**-1)*(15.*mV-v+VT)/
    (exp((15.*mV-v+VT)/(5.*mV))-1.)/ms*(1.-n)-.5*exp((10.*mV-v+VT)/(40.*mV))/ms*n : 1
dh/dt = 0.128*exp((17.*mV-v+VT)/(18.*mV))/ms*(1.-h)-4./(1+exp((40.*mV-v+VT)/(5.*mV)))/ms*h : 1
I : amp
''')

hodgkin_huxley_eqs_with_synaptic_conductance = br.Equations('''
dv/dt = (gl*(El-v) - g_na*(m*m*m)*h*(v-ENa) - g_kd*(n*n*n*n)*(v-EK) + I + ih - g_syn*(v-E_i))/Cm : volt
dm/dt = 0.32*(mV**-1)*(13.*mV-v+VT)/
    (exp((13.*mV-v+VT)/(4.*mV))-1.)/ms*(1-m)-0.28*(mV**-1)*(v-VT-40.*mV)/
    (exp((v-VT-40.*mV)/(5.*mV))-1.)/ms*m : 1
dn/dt = 0.032*(mV**-1)*(15.*mV-v+VT)/
    (exp((15.*mV-v+VT)/(5.*mV))-1.)/ms*(1.-n)-.5*exp((10.*mV-v+VT)/(40.*mV))/ms*n : 1
dh/dt = 0.128*exp((17.*mV-v+VT)/(18.*mV))/ms*(1.-h)-4./(1+exp((40.*mV-v+VT)/(5.*mV)))/ms*h : 1
I : amp
g_syn: siemens
''')

# # First H-current from Izhikevich p. 48
# ghbar = 40. * nS #Other values would be 0.5, 2, 3.5, 20 depending on neuron type (Rothman, Manis)
# Eh = -43*mV
# V_half_h = -75.0
# k_h = -5.5
# V_max_h = -75.
# sigma_h = 15.
# C_amp_h = 1000.
# C_base_h = 100.
# eqs_ih = """
# ih = ghbar*r*(Eh-v) : amp
# dr/dt= (rinf-r)/rtau : 1
# rinf = 1. / (1+exp((V_half_h - v/mV) / k_h)) : 1
# rtau = (C_base_h + C_amp_h * exp(-(V_max_h-v/mV)**2./sigma_h**2)) * ms : second
# """
# Second H-current from Izhikevich p. 48


ih_params = {
    "ghbar": 60. * nS,  # Other values would be 0.5, 2, 3.5, 20 depending on neuron type (Rothman, Manis)
    "Eh": -1. * mV,
    "V_half_h": -98.0,
    "k_h": -5.5,
    "V_max_h": -75.,
    "sigma_h": 25.,
    "C_amp_h": 60.,
    "C_base_h": 40.,
}

eqs_ih = """
ih = ghbar*r*(Eh-v) : amp
dr/dt= (rinf-r)/rtau : 1
rinf = 1. / (1+exp((V_half_h - v/mV) / k_h)) : 1
rtau = 0.1*(C_base_h + C_amp_h * exp(-(V_max_h-v/mV)**2./sigma_h**2)) * ms : second
"""

# eqs_ih = """
# ih = ghbar*r*(Eh-v) : amp
# dr/dt=(rinf-r)/rtau : 1
# rinf = 1. / (1+exp((v/mV + 90.) / 7.)) : 1
# # rtau = ((100000. / (237.*exp((v/mV+60.) / 12.) + 17.*exp(-(v/mV+60.) / 14.))) + 25.)*ms : second
# rtau = 0.01*(100. + 1000. * exp(-(-76.-v/mV)**2./15.**2)) * ms : second #From Izhikevich
# """


delta_synapse_model = 'synaptic_strength: volt'
delta_synapse_on_pre = 'v+=synaptic_strength'
delta_synapse_param = {}

exponential_synapse = """
dg/dt = -g/tau_syn : siemens
g_syn_post = g :siemens (summed)
tau_syn : second
synaptic_strength : siemens
"""
exponential_synapse_on_pre = "g+=synaptic_strength"
exponential_synapse_params = {
    "tau_syn": 5 * ms
}

spike_threshold = +40 * mV

## With voltage delta synapse
neuron_eqs = hodgkin_huxley_eqs+eqs_ih

synapse_model = delta_synapse_model
synapse_on_pre = delta_synapse_on_pre
synapse_params = delta_synapse_param

neuron_params = hodgkin_huxley_params
neuron_params.update(ih_params)

inhibition_off = 0.0 * mV
inhibition_on = -30*mV

# ### 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(neuron_params)
network_params.update(synapse_params)

record_variables = ['v', 'ih']
integration_method = 'exponential_euler'

initial_states = {
    "v": hodgkin_huxley_params["El"]
}

threshold_eqs = """
v_threshold: volt
"""

neuron = br.NeuronGroup(N=1, \
                        model=neuron_eqs + threshold_eqs, \
                        threshold='v > v_threshold', \
                        refractory='v > v_threshold', \
                        method=integration_method)

neuron.v_threshold = spike_threshold
set_parameters_from_dict(neuron, initial_states)
neuron.I = 0.5 * nA

spike_recorder = br.SpikeMonitor(source=neuron)
neuron_state_recorder = br.StateMonitor(neuron, record_variables, record=True)

auto_synapse = br.Synapses(source=neuron, target=neuron, model=synapse_model, on_pre=synapse_on_pre,
                           delay=0.0 * ms)
auto_synapse.connect()
set_parameters_from_dict(auto_synapse, synapse_params)
net = br.Network(neuron)
net.add(auto_synapse)
net.add(spike_recorder)
net.add(neuron_state_recorder)
net.store()


def run_sim(delay=0.0, pulse_strength=0.0, record_states=False, duration=50 * ms):
    net.restore()
    neuron_state_recorder.record = record_states
    auto_synapse.delay = delay
    auto_synapse.synaptic_strength = pulse_strength

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


run_sim(delay=3 * ms, record_states=True, pulse_strength=inhibition_off, duration=200 * ms)

t_state = neuron_state_recorder.t
v_state = neuron_state_recorder[1].v
ih_state = neuron_state_recorder[1].ih
spike_train = spike_recorder.spike_trains()[0]

fig, ax1 = plt.subplots()
ax1.plot(t_state / ms, v_state / ms, 'C2')

plt.show()

period_unperturbed = get_mean_period(spike_train)
print(period_unperturbed)

n_simulations = 40
inhibitory_delay = np.linspace(1 * ms, period_unperturbed - 1 * ms, n_simulations)

spike_trains_dict = {}
periods = []


for delay in inhibitory_delay:
    run_sim(delay, inhibition_on, record_states=False, duration=100 * ms)
    spike_train = spike_recorder.spike_trains()[0]
    #     print(spike_train)
    spike_trains_dict[delay / ms] = spike_train  # Careful, delay is not a Quantity anymore because in must be hashable
    periods.append(get_mean_period(spike_train))

periods = np.array(periods)

plt.plot(inhibitory_delay/ms, periods/ms)
plt.hlines(period_unperturbed/ms, inhibitory_delay[0]/ms, inhibitory_delay[-1]/ms, label="No inhibition")
plt.xlabel("Delay spike to inhibition (ms)")
plt.ylabel("Period (ms")
plt.legend()


t_diff_dict = get_difference_in_periods(spike_trains_dict, period_unperturbed)


phases = np.array([delay / period_unperturbed for delay in inhibitory_delay])
phase_diff = np.array([-t_diff_dict[delay / ms] / period_unperturbed for delay in inhibitory_delay])

optimal_inhibitory_delay_idx = np.argmax(phase_diff)
optimal_inhibitory_delay = list(t_diff_dict.keys())[optimal_inhibitory_delay_idx]  # assumes a single maximum
optimal_inhibitory_delay_phase = phases[optimal_inhibitory_delay_idx]  # assumes a single maximum



plt.figure()
plt.plot(phases, phase_diff)
plt.vlines(optimal_inhibitory_delay, -1, 1)
plt.title("Phase response curve")
plt.ylim(-1, 1)

plt.figure()

if optimal_inhibitory_delay_idx == 0:
    delta_i = phases
else:
    delta_i = phases[:-optimal_inhibitory_delay_idx]

print(delta_i.shape)

delta_iplus1 = delta_i + 1 + phase_diff[optimal_inhibitory_delay_idx:] - phase_diff[optimal_inhibitory_delay_idx]


plt.plot(delta_i, delta_iplus1)
plt.plot(delta_i, delta_i, '--')
plt.hlines(1, 0, 1)
plt.xlabel("Delta i")
plt.ylabel("Delta i+1")
plt.title("Iterative map")
plt.ylim(0, 1.2)
plt.xlim(0, 1)

# plot_mean_response(t_diff_dict)


# plot_phase_response_curve(t_diff_dict, t_isi)


# plot_spiking(spike_trains_dict)

plt.show()