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


rate = 20.0
trials = 20
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 plot_count_fano(ax1, ax2, spikes):
    wins = np.logspace(-2, 0.0, 200)
    mean_counts = np.zeros(len(wins))
    var_counts = np.zeros(len(wins))
    for k, win in enumerate(wins):
        counts = []
        for times in spikes:
            c, _ = np.histogram(times, np.arange(0.0, duration, win))
            counts.extend(c)
        mean_counts[k] = np.mean(counts)
        var_counts[k] = np.var(counts)
    ax1.plot(mean_counts, var_counts, **lsA)
    ax1.set_xlabel('Mean count')
    ax1.set_xlim(0.0, 20.0)
    ax1.set_ylim(0.0, 20.0)
    ax2.plot(1000.0*wins, var_counts/mean_counts, **lsB)
    ax2.set_xlabel('Window', 'ms')
    ax2.set_ylim(0.0, 1.2)
    ax2.set_xscale('log')
    ax2.set_xticks([10, 100, 1000])
    ax2.set_xticklabels(['10', '100', '1000'])

    
if __name__ == "__main__":
    homspikes = hompoisson(rate, trials, duration)
    inhspikes = oupifspikes(rate, trials, duration, dt, 0.3, drate, tau)
    fig, axs = plt.subplots(2, 2)
    fig.subplots_adjust(**adjust_fs(fig, top=0.5, right=1.5))
    plot_count_fano(axs[0,0], axs[0,1], homspikes)
    axs[0,0].text(0.1, 0.95, 'Poisson', transform=axs[0,0].transAxes)
    axs[0,0].set_xlabel('')
    axs[0,1].set_xlabel('')
    axs[0,0].xaxis.set_major_formatter(mpt.NullFormatter())
    axs[0,1].xaxis.set_major_formatter(mpt.NullFormatter())
    plot_count_fano(axs[1,0], axs[1,1], inhspikes)
    axs[1,0].text(0.1, 0.95, 'OU noise', transform=axs[1,0].transAxes)
    fig.text(0.01, 0.58, 'Count variance', va='center', rotation='vertical')
    fig.text(0.53, 0.58, 'Fano factor', va='center', rotation='vertical')
    plt.savefig('fanoexamples.pdf')
    plt.close()