"""Original Gaussian gate-alignment model; write the PNG in the working directory."""
from pathlib import Path
from math import erf, log, pi, sqrt

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter
import numpy as np

# Both offset and window width are measured in power-FWHM units.
def capture(offset, width):
    a = 2 * sqrt(log(2))
    return 0.5 * (erf(a * (offset + width / 2)) - erf(a * (offset - width / 2)))

def density(time):
    return sqrt(4 * log(2) / pi) * np.exp(-4 * log(2) * time**2)

assert np.isclose(capture(0.0, 1.0), 0.760968108550488)
assert np.isclose(capture(0.5, 1.0), 0.4907341611244005)
assert np.isclose(capture(0.0, 2.0), 0.981468322248801)
lo, hi = 0.0, 5.0
for _ in range(90):
    mid = (lo + hi) / 2
    if capture(0.0, mid) < 0.99:
        lo = mid
    else:
        hi = mid
width99 = (lo + hi) / 2
assert np.isclose(width99, 2.1877079813247855, rtol=1e-12)
for width in [0.5, 1.0, 2.0, width99]:
    for offset in [-1.5, -0.5, 0.0, 0.5, 1.5]:
        x = np.linspace(offset - width / 2, offset + width / 2, 60001)
        assert np.isclose(np.trapezoid(density(x), x), capture(offset, width), rtol=1e-8, atol=1e-10)
        assert np.isclose(capture(offset, width), capture(-offset, width), atol=1e-14)
    d = np.linspace(0.001, 2.0, 501)
    derivative = density(d + width / 2) - density(d - width / 2)
    assert np.all(derivative < 0)
    numerical = np.array([(capture(t + 1e-5, width) - capture(t - 1e-5, width)) / 2e-5 for t in d])
    assert np.allclose(derivative, numerical, atol=1e-8, rtol=1e-6)

sigma = 2e-9
energy = 1e-9
power_fwhm = 2 * sqrt(2 * log(2)) * sigma
peak = energy / (sqrt(2 * pi) * sigma)
assert np.isclose(power_fwhm / 1e-9, 4.709640090061899)
assert np.isclose(peak, 0.19947114020071635)
assert np.isclose(power_fwhm * sqrt(2) / power_fwhm, sqrt(2))

offsets = np.linspace(-2.1, 2.1, 1001)
fig, ax = plt.subplots(figsize=(9.0, 4.6), dpi=100, facecolor="white")
for width, color in [(1.0, "#0072B2"), (2.0, "#D55E00")]:
    fraction = np.array([capture(d, width) for d in offsets])
    assert np.all((fraction >= 0) & (fraction <= 1))
    ax.plot(offsets, fraction, lw=2.4, color=color,
            label=f"Gate width {width:g}τ; centred capture {capture(0, width):.2%}")
    ax.plot(0, capture(0, width), "o", color=color, ms=5)
ax.plot(0.5, capture(0.5, 1), "o", color="#222222", ms=5)
ax.annotate("τ-wide gate with offset τ/2:\n49.07% captured",
            xy=(0.5, capture(0.5, 1)), xytext=(1.05, 0.76),
            fontsize=9, arrowprops=dict(arrowstyle="->", color="#444444"))
ax.axvline(0, color="#777777", ls="--", lw=0.8, zorder=0)
ax.axhline(1, color="#777777", ls=":", lw=0.8, zorder=0)
ax.set(xlim=(-2.1, 2.1), ylim=(-0.015, 1.06),
       xlabel="Gate-centre offset d / τ", ylabel="Captured pulse energy Ug / U")
ax.yaxis.set_major_formatter(PercentFormatter(1.0))
ax.set_facecolor("white")
ax.spines[["top", "right"]].set_visible(False)
ax.grid(alpha=0.16)
ax.legend(loc="lower center", frameon=False, fontsize=9)
fig.suptitle("Gaussian pulse capture depends on gate width and alignment", fontsize=13, y=0.97)
fig.text(0.5, 0.89, "τ is the power FWHM; ideal integration of incident optical power",
         ha="center", fontsize=10)
fig.subplots_adjust(left=0.09, right=0.98, bottom=0.16, top=0.82)
fig.savefig(Path.cwd() / "gaussian-pulse-gate-alignment.png", facecolor="white")
plt.close(fig)