import numpy as np
import matplotlib.pyplot as plt
from plotstyle import *


rate = 20.0
trials = 10
duration = 100.0
dt = 0.001
drate = 50.0
tau = 0.1;


def hompoisson(rate, trials, duration) :
    spikes = []
    for k in range(trials) :
        times = []
        t = 0.0
        while t < duration :
            t += np.random.exponential(1/rate)
            times.append( t )
        spikes.append( times )
    return spikes


def inhompoisson(rate, trials, dt) :
    spikes = []
    p = rate*dt
    for k in range(trials) :
        x = np.random.rand(len(rate))
        times = dt*np.nonzero(x<p)[0]
        spikes.append( times )
    return spikes


def pifspikes(input, trials, dt, D=0.1) :
    vreset = 0.0
    vthresh = 1.0
    tau = 1.0
    spikes = []
    for k in range(trials) :
        times = []
        v = vreset
        noise = np.sqrt(2.0*D)*np.random.randn(len(input))/np.sqrt(dt)
        for k in range(len(noise)) :
            v += (input[k]+noise[k])*dt/tau
            if v >= vthresh :
                v = vreset
                times.append(k*dt)
        spikes.append( times )
    return spikes


def oupifspikes(rate, trials, duration, dt, D, drate, tau):
    # OU noise:
    rng = np.random.RandomState(54637281)
    time = np.arange(0.0, duration, dt)
    x = np.zeros(time.shape)+rate
    n = rng.randn(len(time))*drate*tau/np.sqrt(dt) + rate
    for k in range(1,len(x)) :
        x[k] = x[k-1] + (n[k]-x[k-1])*dt/tau
    x[x<0.0] = 0.0
    spikes = pifspikes(x, trials, dt, D)
    return spikes


def isis( spikes ) :
    isi = []
    for k in range(len(spikes)) :
        isi.extend(np.diff(spikes[k]))
    return isi


def plotisih( ax, isis, binwidth=None ) :
    if binwidth == None :
        nperbin = 200.0    # average number of isis per bin
        bins = len(isis)/nperbin  # number of bins
        binwidth = np.max(isis)/bins
        if binwidth < 5e-4 :     # half a millisecond
            binwidth = 5e-4
    h, b = np.histogram(isis, np.arange(0.0, np.max(isis)+binwidth, binwidth), density=True)
    ax.text(0.9, 0.85, 'rate={:.0f}Hz'.format(1.0/np.mean(isis)), ha='right', transform=ax.transAxes)
    ax.text(0.9, 0.7, 'mean={:.0f}ms'.format(1000.0*np.mean(isis)), ha='right', transform=ax.transAxes)
    ax.text(0.9, 0.55, 'CV={:.2f}'.format(np.std(isis)/np.mean(isis)), ha='right', transform=ax.transAxes)
    ax.set_xlabel('ISI', 'ms')
    ax.set_ylabel('p(ISI)', '1/s')
    ax.bar( 1000.0*b[:-1], h, bar_fac*1000.0*np.diff(b), **fsA)

    
def plot_hom_isih(ax):
    homspikes = hompoisson(rate, trials, duration)
    ax.set_xlim(0.0, 150.0)
    ax.set_ylim(0.0, 31.0)
    ax.set_xticks(np.arange(0.0, 151.0, 50.0))
    ax.set_yticks(np.arange(0.0, 31.0, 10.0))
    tt = np.linspace(0.0, 0.15, 100)
    ax.plot(1000.0*tt, rate*np.exp(-rate*tt), **lsB)
    plotisih(ax, isis(homspikes), 0.005)

    
def plot_inhom_isih(ax):
    inhspikes = oupifspikes(rate, trials, duration, dt, 0.3, drate, tau)
    ax.set_xlim(0.0, 150.0)
    ax.set_ylim(0.0, 31.0)
    ax.set_xticks(np.arange(0.0, 151.0, 50.0))
    ax.set_yticks(np.arange(0.0, 31.0, 10.0))
    plotisih(ax, isis(inhspikes), 0.005)

    
if __name__ == "__main__":
    fig, (ax1, ax2) = plt.subplots(1, 2)
    fig.subplots_adjust(**adjust_fs(fig, top=0.5, right=1.5))
    plot_hom_isih(ax1)
    plot_inhom_isih(ax2)
    plt.savefig('isihexamples.pdf')
    plt.close()