### import ####################################################################


import matplotlib.pyplot as plt
plt.close('all')

import numpy as np

import WrightTools as wt


### define ####################################################################


def get_signal(d, tau, pulsewidth=10, offset=0):
    # pulse
    pulse = np.exp((-d**2)/(pulsewidth**2))
    # signal
    sig = np.zeros(d.shape)
    sig[d<=0] = np.exp(d[d<=0]/tau)
    sig[d<=0] += offset
    sig /= sig.max()
    # finish
    #sig = np.convolve(sig, pulse, mode='same')
    return sig


def logarithmic_stepping(p_tau, p_npts, n_tau, n_npts):
    # positive
    p_xi = np.arange(0, p_npts)
    p_delays = p_tau * np.log((p_xi.size+1)/(p_xi+1))
    # negative
    n_xi = np.arange(0, n_npts)
    n_delays = -n_tau * np.log((n_xi.size+1)/(n_xi+1))
    return np.hstack((n_delays, [0], p_delays))


tau = 200


d = logarithmic_stepping(50, 3, 200, 15)


### workspace #################################################################


if True:
    fig, gs = wt.artists.create_figure(width=13, cols=[1, 1], nrows=1)
    # delay space
    ax = plt.subplot(gs[0, 0]) 
    ds = np.linspace(-1500, 1500, 1000)
    sig = get_signal(ds, tau)
    plt.plot(ds, sig, c='b', lw=2, alpha=0.5)
    sig = get_signal(ds, tau, offset=0.5)
    plt.plot(ds, sig, c='r', lw=2, alpha=0.5)
    sig = get_signal(ds, tau*2, offset=0)
    plt.plot(ds, sig, c='g', lw=2, alpha=0.5)
    plt.xlim(-1250, 100)
    plt.ylim(-0.1, 1.1)
    for x in d:
        plt.axvline(x, c='k', zorder=0)
    plt.axvline(0, lw=3, c='k')
    ax.set_xlabel('delay', fontsize=18)
    ax.set_ylabel('signal', fontsize=18)
    plt.grid(ls=':')
    # index space
    ax = plt.subplot(gs[0, 1])
    d = logarithmic_stepping(50, 3, 200, 15)
    sig = get_signal(d, tau)
    plt.scatter(np.arange(sig.size), sig, c='b', edgecolor='none', s=50, alpha=0.5)
    sig = get_signal(d, tau, offset=0.5)
    plt.scatter(np.arange(sig.size), sig, c='r', edgecolor='none', s=50, alpha=0.5)
    sig = get_signal(d, tau*2, offset=0)
    plt.scatter(np.arange(sig.size), sig, c='g', edgecolor='none', s=50, alpha=0.5) 
    i = np.argmin(np.abs(d))
    plt.axvline(i, lw=3, c='k')
    plt.grid(ls=':')
    plt.ylim(-0.1, 1.1)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.set_xlim(0-1, sig.size)
    ax.set_xlabel('index', fontsize=18)
    # finish
    wt.artists.savefig('exponential.png')