import matplotlib.pyplot as plt
import numpy as np
from IPython import embed

figsize=(6,3)

def set_rc():
    plt.rcParams['xtick.major.size'] = 5
    plt.rcParams['xtick.minor.size'] = 5
    plt.rcParams['xtick.major.width'] = 2
    plt.rcParams['xtick.minor.width'] = 2
    plt.rcParams['ytick.major.size'] = 5
    plt.rcParams['ytick.minor.size'] = 5
    plt.rcParams['ytick.major.width'] = 2
    plt.rcParams['ytick.minor.width'] = 2
    plt.rcParams['xtick.direction'] = "out"
    plt.rcParams['ytick.direction'] = "out"


def create_spikes(nspikes=11, duration=0.5, seed=1000):
    rng = np.random.RandomState(seed)
    x = np.linspace(0.0, 1.0, nspikes)
    # double gaussian rate profile:
    rate = np.exp(-0.5*((x-0.35)/0.25)**2.0)
    rate += 1.*np.exp(-0.5*((x-0.9)/0.05)**2.0)
    isis = 1.0/rate
    isis += rng.randn(len(isis))*0.2
    times = np.cumsum(isis)
    times *= 1.05*duration/times[-1]
    times += 0.01
    return times


def gaussian(sigma, dt):
    x = np.arange(-4*sigma, 4*sigma, dt)
    y = np.exp(-0.5 * (x / sigma)**2)/np.sqrt(2*np.pi)/sigma; 
    return x, y

    
def setup_axis(spikes_ax, rate_ax):
    spikes_ax.spines["left"].set_visible(False)
    spikes_ax.spines["right"].set_visible(False)
    spikes_ax.spines["top"].set_visible(False)
    spikes_ax.yaxis.set_ticks_position('left')
    spikes_ax.xaxis.set_ticks_position('bottom')
    spikes_ax.set_yticks([])
    spikes_ax.set_ylim(-0.2, 1.0)
    #spikes_ax.set_ylabel("Spikes")
    spikes_ax.text(-0.1, 0.5, "Spikes", transform=spikes_ax.transAxes, rotation='vertical', va='center')
    #spikes_ax.text(-0.125, 1.2, "A", transform=spikes_ax.transAxes)
    spikes_ax.set_xlim(-1, 500)
    #spikes_ax.set_xticklabels(np.arange(0., 600, 100))

    rate_ax.spines["right"].set_visible(False)
    rate_ax.spines["top"].set_visible(False)
    rate_ax.yaxis.set_ticks_position('left')
    rate_ax.xaxis.set_ticks_position('bottom')
    rate_ax.set_xlabel('Time [ms]')
    #rate_ax.set_ylabel('Firing rate [Hz]')
    rate_ax.text(-0.1, 0.5, "Rate [Hz]", transform=rate_ax.transAxes, rotation='vertical', va='center')
    #rate_ax.text(-0.125, 1.15, "B", transform=rate_ax.transAxes)   
    rate_ax.set_xlim(0, 500)
    #rate_ax.set_xticklabels(np.arange(0., 600, 100))
    rate_ax.set_ylim(0, 60)
    rate_ax.set_yticks(np.arange(0,65,20))


def plot_bin_method():
    dt = 1e-5
    duration = 0.5
    
    spike_times = create_spikes()
    t = np.arange(0., duration, dt)
    
    bins = np.arange(0, 0.55, 0.05)
    count, _ = np.histogram(spike_times, bins)
    
    plt.xkcd()
    set_rc()
    fig = plt.figure()
    fig.set_size_inches(*figsize)
    fig.set_facecolor('white')
    spikes = plt.subplot2grid((7,1), (0,0), rowspan=3, colspan=1)
    rate_ax = plt.subplot2grid((7,1), (3,0), rowspan=4, colspan=1)
    setup_axis(spikes, rate_ax)

    for ti in spike_times:
        ti *= 1000.0
        spikes.plot([ti, ti], [0., 1.], '-b', lw=2)

    #spikes.vlines(1000.0*spike_times, 0., 1., color="darkblue", lw=1.25)
    for tb in 1000.0*bins :
        spikes.plot([tb, tb], [-2.0, 0.75], '-', color="#777777", lw=1, clip_on=False)
    #spikes.vlines(1000.0*np.hstack((0,bins)), -2.0, 1.25, color="#777777", lw=1, linestyles='-', clip_on=False)
    for i,c in enumerate(count):
        spikes.text(1000.0*(bins[i]+0.5*bins[1]), 1.1, str(c), color='#CC0000', ha='center')

    rate = count / 0.05
    rate_ax.step(1000.0*bins, np.hstack((rate, rate[-1])), color='#FF9900', where='post')
    fig.tight_layout()
    fig.savefig("binmethod.pdf")
    plt.close()


def plot_conv_method():
    dt = 1e-5
    duration = 0.5
    spike_times = create_spikes() 
    kernel_time, kernel = gaussian(0.015, dt)
    
    t = np.arange(0., duration, dt)
    rate = np.zeros(t.shape)
    rate[np.asarray(np.round(spike_times[:-1]/dt), dtype=int)] = 1  
    rate = np.convolve(rate, kernel, mode='same')
    rate = np.roll(rate, -1)

    plt.xkcd()
    set_rc()
    fig = plt.figure()
    fig.set_size_inches(*figsize)
    fig.set_facecolor('white')
    spikes = plt.subplot2grid((7,1), (0,0), rowspan=3, colspan=1)
    rate_ax = plt.subplot2grid((7,1), (3,0), rowspan=4, colspan=1)
    setup_axis(spikes, rate_ax)
    
    #spikes.vlines(1000.0*spike_times, 0., 1., color="darkblue", lw=1.5, zorder=2)
    for ti in spike_times:
        ti *= 1000.0
        spikes.plot([ti, ti], [0., 1.], '-b', lw=2)
        spikes.plot(1000*kernel_time + ti, kernel/np.max(kernel), color='#cc0000', lw=1, zorder=1)

    rate_ax.plot(1000.0*t, rate, color='#FF9900', lw=2, zorder=2)
    rate_ax.fill_between(1000.0*t, rate, np.zeros(len(rate)), color='#FFFF66')
    #rate_ax.fill_between(t, rate, np.zeros(len(rate)), color="red", alpha=0.5)
    #rate_ax.set_ylim([0, 50])
    #rate_ax.set_yticks(np.arange(0,75,25))
    fig.tight_layout()
    fig.savefig("convmethod.pdf")


def plot_isi_method():
    spike_times = create_spikes()
    
    plt.xkcd()
    set_rc()
    fig = plt.figure()
    fig.set_size_inches(*figsize)
    fig.set_facecolor('white')
    
    spikes = plt.subplot2grid((7,1), (0,0), rowspan=3, colspan=1)
    rate = plt.subplot2grid((7,1), (3,0), rowspan=4, colspan=1)
    setup_axis(spikes, rate)
    
    spike_times = np.hstack((0.005, spike_times))
    #spikes.vlines(1000*spike_times, 0., 1., color="blue", lw=2)
    for i in range(1, len(spike_times)):
        t_start = 1000*spike_times[i-1]
        t = 1000*spike_times[i]
        spikes.plot([t_start, t_start], [0., 1.], '-b', lw=2)
        spikes.annotate(s='', xy=(t_start, 0.5), xytext=(t,0.5), arrowprops=dict(arrowstyle='<->'), color='red')
        spikes.text(0.5*(t_start+t), 0.75, 
                    "{0:.0f}".format((t - t_start)),
                    color='#CC0000', ha='center')

    #spike_times = np.hstack((0, spike_times))
    i_rate = 1./np.diff(spike_times)
    rate.step(1000*spike_times, np.hstack((i_rate, i_rate[-1])),color='#FF9900', lw=2, where="post")
    
    fig.tight_layout()
    fig.savefig("isimethod.pdf")
    
    
if __name__ == '__main__':
    plot_isi_method()
    plot_conv_method()
    plot_bin_method()