import numpy as np
import matplotlib.pyplot as plt

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

time = np.arange(-0.05, 0.05, 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(figsize = (20/inch_factor, 12/inch_factor))
ax1 = fig.add_subplot(211)
plt.yticks(fontsize=18)
ax2 = fig.add_subplot(212, sharex=ax1)
plt.setp(ax1.get_xticklabels(), visible=False)
ax1.plot(time*1000, signal, color = 'royalblue', lw = 1)
ax2.plot(time*1000, freq, color = 'royalblue', lw = 3)

ax1.set_ylabel("field [mV]", fontsize = 22)


ax1.yaxis.set_label_coords(-0.1, 0.5)
ax2.set_xlabel("time [ms]", fontsize = 22)
ax2.set_ylabel("frequency [Hz]", fontsize = 22)
ax2.yaxis.set_label_coords(-0.1, 0.5)
plt.xticks(fontsize=18)
plt.yticks(fontsize=18)
fig.tight_layout()
plt.show()
#plt.savefig('stimulus_chirp.png')