# -*- coding: utf-8 -*-
#
# CampbellSiegert.py
#
# This file is part of NEST.
#
# Copyright (C) 2004 The NEST Initiative
#
# NEST is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# NEST is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with NEST.  If not, see <http://www.gnu.org/licenses/>.

"""
Campbell & Siegert approximation example
----------------------------------------

Example script that applies Campbell's theorem and Siegert's rate
approximation to and integrate-and-fire neuron.

This script calculates the firing rate of an integrate-and-fire neuron
in response to a series of Poisson generators, each specified with a
rate and a synaptic weight. The calculated rate is compared with a
simulation using the ``iaf_psc_alpha`` model



References
~~~~~~~~~~

 .. [1] Papoulis A (1991). Probability, Random Variables, and
        Stochastic Processes, McGraw-Hill
 .. [2] Siegert AJ (1951). On the first passage time probability problem,
        Phys Rev 81: 617-623

Authors
~~~~~~~

S. Schrader, Siegert implementation by T. Tetzlaff
"""

###############################################################################
# First, we import all necessary modules for simulation and analysis. Scipy
# should be imported before nest.

import nest
import numpy as np
from scipy.optimize import fmin
from scipy.special import erf

###############################################################################
# We first set the parameters of neurons, noise and the simulation. First
# settings are with a single Poisson source, second is with two Poisson
# sources with half the rate of the single source. Both should lead to the
# same results.

weights = [0.1]  # (mV) psp amplitudes
rates = [10000.0]  # (1/s) rate of Poisson sources
# weights = [0.1, 0.1]    # (mV) psp amplitudes
# rates = [5000., 5000.]  # (1/s) rate of Poisson sources

C_m = 250.0  # (pF) capacitance
E_L = -70.0  # (mV) resting potential
I_e = 0.0  # (nA) external current
V_reset = -70.0  # (mV) reset potential
V_th = -55.0  # (mV) firing threshold
t_ref = 2.0  # (ms) refractory period
tau_m = 10.0  # (ms) membrane time constant
tau_syn_ex = 0.5  # (ms) excitatory synaptic time constant
tau_syn_in = 2.0  # (ms) inhibitory synaptic time constant

simtime = 20000  # (ms) duration of simulation
n_neurons = 10  # number of simulated neurons

###############################################################################
# For convenience we define some units.

pF = 1e-12
ms = 1e-3
pA = 1e-12
mV = 1e-3

mu = 0.0
sigma2 = 0.0
J = []

assert len(weights) == len(rates)

###############################################################################
# In the following we analytically compute the firing rate of the neuron
# based on Campbell's theorem [1]_ and Siegerts approximation [2]_.

for rate, weight in zip(rates, weights):
    if weight > 0:
        tau_syn = tau_syn_ex
    else:
        tau_syn = tau_syn_in

    # We define the form of a single PSP, which allows us to match the
    # maximal value to or chosen weight.

    def psp(x):
        return -(
            (C_m * pF)
            / (tau_syn * ms)
            * (1 / (C_m * pF))
            * (np.exp(1) / (tau_syn * ms))
            * (
                ((-x * np.exp(-x / (tau_syn * ms))) / (1 / (tau_syn * ms) - 1 / (tau_m * ms)))
                + (np.exp(-x / (tau_m * ms)) - np.exp(-x / (tau_syn * ms)))
                / ((1 / (tau_syn * ms) - 1 / (tau_m * ms)) ** 2)
            )
        )

    min_result = fmin(psp, [0], full_output=1, disp=0)

    # We need to calculate the PSC amplitude (i.e., the weight we set in NEST)
    # from the PSP amplitude, that we have specified above.

    fudge = -1.0 / min_result[1]
    J.append(C_m * weight / (tau_syn) * fudge)

    # We now use Campbell's theorem to calculate mean and variance of
    # the input due to the Poisson sources. The mean and variance add up
    # for each Poisson source.

    mu += rate * (J[-1] * pA) * (tau_syn * ms) * np.exp(1) * (tau_m * ms) / (C_m * pF)

    sigma2 += (
        rate
        * (2 * tau_m * ms + tau_syn * ms)
        * (J[-1] * pA * tau_syn * ms * np.exp(1) * tau_m * ms / (2 * (C_m * pF) * (tau_m * ms + tau_syn * ms))) ** 2
    )

mu += E_L * mV
sigma = np.sqrt(sigma2)

###############################################################################
# Having calculate mean and variance of the input, we can now employ
# Siegert's rate approximation.

num_iterations = 100
upper = (V_th * mV - mu) / (sigma * np.sqrt(2))
lower = (E_L * mV - mu) / (sigma * np.sqrt(2))
interval = (upper - lower) / num_iterations
tmpsum = 0.0
for cu in range(0, num_iterations + 1):
    u = lower + cu * interval
    f = np.exp(u**2) * (1 + erf(u))
    tmpsum += interval * np.sqrt(np.pi) * f
r = 1.0 / (t_ref * ms + tau_m * ms * tmpsum)

###############################################################################
# We now simulate neurons receiving Poisson spike trains as input,
# and compare the theoretical result to the empirical value.

nest.ResetKernel()

nest.set_verbosity("M_WARNING")
neurondict = {
    "V_th": V_th,
    "tau_m": tau_m,
    "tau_syn_ex": tau_syn_ex,
    "tau_syn_in": tau_syn_in,
    "C_m": C_m,
    "E_L": E_L,
    "t_ref": t_ref,
    "V_m": E_L,
    "V_reset": E_L,
}

###############################################################################
# Neurons and devices are instantiated. We set a high threshold as we want
# free membrane potential. In addition we choose a small resolution for
# recording the membrane to collect good statistics.

n = nest.Create("iaf_psc_alpha", n_neurons, params=neurondict)
n_free = nest.Create("iaf_psc_alpha", params=dict(neurondict, V_th=1e12))
pg = nest.Create("poisson_generator", len(rates), {"rate": rates})
vm = nest.Create("voltmeter", params={"interval": 0.1})
sr = nest.Create("spike_recorder")

###############################################################################
# We connect devices and neurons and start the simulation.

pg_n_synspec = {"weight": np.tile(J, ((n_neurons), 1)), "delay": 0.1}
nest.Connect(pg, n, syn_spec=pg_n_synspec)
nest.Connect(pg, n_free, syn_spec={"weight": [J]})
nest.Connect(vm, n_free)
nest.Connect(n, sr)

nest.Simulate(simtime)

###############################################################################
# Here we read out the recorded membrane potential. The first 500 steps are
# omitted so initial transients do not perturb our results. We then print the
# results from theory and simulation.

v_free = vm.events["V_m"]
Nskip = 500
print(f"mean membrane potential (actual / calculated): {np.mean(v_free[Nskip:])} / {mu * 1000}")
print(f"variance (actual / calculated): {np.var(v_free[Nskip:])} / {sigma2 * 1e6}")
print(f"firing rate (actual / calculated): {sr.n_events / (n_neurons * simtime * ms)} / {r}")
