import numpy as np
import matplotlib.pyplot as plt

stimulusrate = 500.  # the eod frequency of the fake fish
currentchirptimes = [0.1]
chirpwidth = 0.05  # ms
chirpsize = 100.
chirpampl = 0.02
chirpkurtosis = 1.
p = 0.
stepsize = 0.00001

time = np.arange(0.0, 0.2, stepsize)
signal = np.zeros(time.shape)
ampl = np.ones(time.shape)
freq = np.ones(time.shape)

ck = 0
csig = 0.5 * chirpwidth / np.power(2.0*np.log(10.0), 0.5/chirpkurtosis)

for k, t in enumerate(time):
    a = 1.
    f = stimulusrate
    if ck < len(currentchirptimes):
        if np.abs(t - currentchirptimes[ck]) < 2.0 * chirpwidth:
            x = t - currentchirptimes[ck]
            g = np.exp(-0.5 * (x/csig)**2)
            f = chirpsize * g + stimulusrate
            a *= 1.0 - chirpampl * g
        elif t > currentchirptimes[ck] + 2.0 * chirpwidth:
            ck += 1
    freq[k] = f
    ampl[k] = a
    p += f * stepsize
    signal[k] = a * np.sin(6.28318530717959 * p)

fig = plt.figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)

ax1.plot(time, signal)
ax2.plot(time, freq)

ax1.set_ylabel("fake fish field [rel]")
ax2.set_xlabel("time [s]")
ax2.set_ylabel("frequency [Hz]")
plt.show()