"""Plot original Gaussian power profiles; output the PNG in the working directory."""
from pathlib import Path
from math import erf

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

energy = 0.50e-6  # J
widths = [10e-9, 40e-9]  # Power FWHM, s; not field-amplitude FWHM.
colors = ["#0072B2", "#D55E00"]
factor = np.sqrt(np.pi / (4 * np.log(2)))
captured_fraction = erf(np.sqrt(np.log(2)))
assert np.isclose(factor, 1.0644670194312262)
assert np.isclose(captured_fraction, 0.760968108550488)
assert np.isclose(energy * 1e6, 0.50)  # W average at 1 MHz.

def power(time, width):
    return energy / (factor * width) * np.exp(-4 * np.log(2) * (time / width) ** 2)

fig, ax = plt.subplots(figsize=(9.0, 4.5), dpi=100, facecolor="white")
time = np.linspace(-100e-9, 100e-9, 5001)
peaks = []
for width, color in zip(widths, colors):
    peak = energy / (factor * width)
    peaks.append(peak)
    wide_time = np.linspace(-10 * max(widths), 10 * max(widths), 100001)
    assert np.isclose(np.trapezoid(power(wide_time, width), wide_time), energy, rtol=1e-10, atol=0.0)
    assert np.allclose(power(np.array([-width / 2, width / 2]), width), peak / 2)
    gate_time = np.linspace(-width / 2, width / 2, 10001)
    gate_energy = np.trapezoid(power(gate_time, width), gate_time)
    assert np.isclose(gate_energy / energy, captured_fraction, rtol=1e-8)
    values = power(time, width)
    ax.plot(time * 1e9, values, color=color, lw=2.3,
            label=f"{width * 1e9:.0f} ns power FWHM; peak {peak:.2f} W")
    ax.fill_between(gate_time * 1e9, power(gate_time, width), color=color, alpha=0.13)
    ax.annotate("", xy=(width * 1e9 / 2, peak / 2), xytext=(-width * 1e9 / 2, peak / 2),
                arrowprops=dict(arrowstyle="<->", color=color, lw=1.3))
    ax.text(width * 1e9 / 2 + 5, peak / 2 + 1, f"{width * 1e9:.0f} ns FWHM",
            color=color, fontsize=10)
assert np.isclose(peaks[0] / peaks[1], 4.0)
ax.set(xlim=(-85, 85), ylim=(-0.7, 53), xlabel="Time from pulse centre t (ns)", ylabel="Radiant power P(t) (W)")
ax.set_facecolor("white")
ax.spines[["top", "right"]].set_visible(False)
ax.grid(alpha=0.17)
ax.legend(loc="upper right", fontsize=9, frameon=False)
fig.suptitle("Same pulse energy, different Gaussian pulse widths", fontsize=14, y=0.97)
fig.text(0.5, 0.89, "U = 0.50 µJ each; shaded power-FWHM windows each contain 76.10% of U",
         ha="center", fontsize=10)
fig.subplots_adjust(left=0.09, right=0.98, bottom=0.15, top=0.82)
fig.savefig(Path.cwd() / "gaussian-pulse-energy-width.png", facecolor="white")
plt.close(fig)