import matplotlib.pyplot as plt
import numpy as np

### supporting functions 
def gauss(x, sig, mu):
    return np.exp(-((x-mu)**2)/(2 * sig**2))

def fn(var, arr):
    return np.argmin((var-arr)**2)

def halfgauss(x, sig, mu):
    g1 = gauss(x, sig, mu)
    edge = np.zeros_like(x)
    edge[fn(mu, x):] += 1
    return g1 * edge

def normalise(arr):
    return arr/arr.max()


### plotting
plot_gap = 1.2

decay = 0.1
fwhm = decay * 2
time = np.linspace(-6*fwhm, 6*fwhm, 1001)

fig, ax = plt.subplots(1, 4, figsize=(6, 6))

colors = ['tab:blue', 'tab:orange', 'tab:green']

for n in range(0, 3):
    radiation = np.zeros_like(time)
    scintillator_response = np.zeros_like(time)
    mcp_response = np.zeros_like(time)
    if n:
        ax[2].plot([-fwhm/(decay), fwhm/(decay)], [0.5+(plot_gap*n),
                   0.5+(plot_gap*n)], 'k-', lw=1)

    if n == 0:
        radiation = gauss(time, decay, -fwhm) + 0.4*gauss(time, decay, fwhm)
        scintillator_response[len(time)//2] += 1
        mcp_response[len(time)//2] += 1

        ax[0].set_title('Radiation', fontsize=10)
        ax[1].set_title('Scintillator Decay', fontsize=10)
        ax[2].set_title('MCP Response', fontsize=10)
        ax[3].set_title('Measured Signal', fontsize=10)

    elif n == 1:
        radiation = gauss(time, decay, -fwhm) + 0.4 * \
            gauss(time, decay, fwhm)
        scintillator_response = halfgauss(time, decay, 0)
        mcp_response = gauss(time, fwhm, 0)

        ax[0].plot([-fwhm/(decay), -fwhm/(decay)],
                   [1+(plot_gap*n), 1.15+(plot_gap*n)], 'k--', lw=1)

        ax[0].plot([fwhm/(decay), fwhm/(decay)],
                   [0.4+(plot_gap*n), 1.15+(plot_gap*n)], 'k--', lw=1)

        ax[0].text(-fwhm/(1.1*decay), 1.05+(plot_gap*n),
                   '$\delta t$', fontsize=9)

    elif n == 2:
        radiation = gauss(time, decay, -fwhm*2) + 0.4 * \
            gauss(time, decay, fwhm*2)
        scintillator_response = halfgauss(time, decay, 0)
        mcp_response = gauss(time, fwhm, 0)
        
        ax[0].plot([-2*fwhm/decay, -2*fwhm/decay], [1+(plot_gap*n),
                   1.15+(plot_gap*n)], 'k--', lw=1)
        
        ax[0].plot([2*fwhm/decay, 2*fwhm/decay],
                   [0.4+(plot_gap*n), 1.15+(plot_gap*n)], 'k--', lw=1)
        
        ax[0].text(-fwhm/(decay), 1.05+(plot_gap*n),
                   '$\delta t$', fontsize=9)

        ax[1].set_xlabel('                                  Time ($\sigma$)')
        ax[0].set_xlabel('Emitted Time ($\sigma $)')
        ax[3].set_xlabel('Recorded Time ($\sigma $)')

    scint_conv = np.convolve(radiation, scintillator_response, 'same')

    ax[0].plot(time/decay,  normalise(radiation /
               np.sum(radiation))+(plot_gap*n))
    ax[1].plot(time/decay, scintillator_response+(plot_gap*n))
    ax[2].plot(time/decay, mcp_response+(plot_gap*n))

    
    delta_func = np.zeros_like(mcp_response)
    delta_func[len(delta_func)//2] += 1
    real_response = np.convolve(
        normalise(scint_conv), mcp_response/np.max(mcp_response), 'same')
    delta_response = normalise(np.convolve(scint_conv, delta_func, 'same'))

    ax[3].plot(time/decay,  real_response/np.max(real_response) +
               (plot_gap*n), color=colors[n])
    ax[3].plot(time/decay,  delta_response +
               (plot_gap*n), color=colors[n], alpha=0.4, ls='--')


for a in ax.ravel():
    a.set_ylim([-0.05, 3.75])

for a in ax[1:].ravel():
    a.set_yticks([])
ax[0].set_yticks([0.5, 1.75, 3.0])
ax[0].set_yticklabels(['Ideal', '$\delta t = 2\\tau_g$',
                      '$\delta t = 4\\tau_g$'], rotation=45, va="center")

# for a in ax[:-1].ravel():
#     a.set_xticks([])
plt.tight_layout()
plt.text(-7.95/decay, 1.6, '$\\ast$', fontsize=14)
plt.text(-4.85/decay, 1.6, '$\\ast$', fontsize=14)
plt.text(-1.75/decay, 1.6, '=', fontsize=14)

plt.text(-11.2/decay, 3.65, 'a)', fontweight='bold')
plt.text(-7.9/decay, 3.65, 'b)', fontweight='bold')
plt.text(-4.8/decay, 3.65, 'c)', fontweight='bold')
plt.text(-1.7/decay, 3.65, 'd)', fontweight='bold')


# plt.savefig('./figure1_generaliseddecay.pdf')
plt.show()