"""Original ideal rectangular-pulse timing diagram; write PNG to the cwd.

Timings are chosen for illustration, not a reconstruction of a published run.
The normalized phase-encoding wavevector includes only applied gradients.
"""
from pathlib import Path

import matplotlib

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


def main():
    duration, first, second, refocus, echo = 2.0, 3.0, 33.0, 30.0, 60.0
    t = np.linspace(-3, 65, 13601)
    gradient = (((t >= first) & (t < first + duration)).astype(float)
                + ((t >= second) & (t < second + duration)).astype(float))
    q = (np.clip((t - first) / duration, 0, 1)
         - np.clip((t - second) / duration, 0, 1))
    integral = np.trapezoid(q**2, t)
    expected = second - first - duration / 3
    assert abs(integral - expected) < 1e-5, (integral, expected)
    fig, axes = plt.subplots(3, 1, figsize=(8, 6), dpi=100, sharex=True,
                             gridspec_kw={"height_ratios": [0.85, 1.2, 1.2]},
                             facecolor="white")
    fig.subplots_adjust(left=0.14, right=0.97, top=0.80, bottom=0.20, hspace=0.3)
    fig.suptitle("Encoding displacement with a refocused gradient pair",
                 fontsize=14, weight="bold", y=0.965, color="#172330")
    fig.text(0.5, 0.905, "Idealized chosen timings; both physical gradient pulses have the same sign",
             ha="center", fontsize=10.5, color="#45535d")
    rf, grad, phase = axes
    rf.hlines(0, -3, 65, color="#7b8992", linewidth=1)
    rf.vlines([0, refocus], [0, 0], [1, 1.8], color="#273c50", linewidth=3)
    rf.text(0, 1.18, "90°", ha="center", fontsize=11)
    rf.text(refocus, 1.98, "180°", ha="center", fontsize=11)
    rf.plot(echo, 0, "o", color="#973c65", markersize=5)
    rf.text(echo, 0.40, "echo", ha="center", fontsize=11, color="#973c65")
    rf.set_ylim(-0.2, 2.5)
    rf.set_yticks([])
    rf.set_ylabel("RF pulses", fontsize=11)
    grad.plot(t, gradient, color="#216c9b", linewidth=2)
    grad.fill_between(t, 0, gradient, color="#216c9b", alpha=0.17)
    grad.annotate("", (second, 1.70), (first, 1.70),
                  arrowprops={"arrowstyle": "<->", "color": "#273c50"})
    grad.text((first + second) / 2, 1.78, r"$\Delta=30\,\mathrm{ms}$",
              ha="center", fontsize=10)
    grad.annotate("", (first + duration, 1.15), (first, 1.15),
                  arrowprops={"arrowstyle": "<->", "color": "#273c50"})
    grad.text(first + duration / 2, 1.26, r"$\delta=2\,\mathrm{ms}$",
              ha="center", fontsize=10)
    grad.set_ylim(-0.15, 2.15)
    grad.set_yticks([0, 1], ["0", "g"])
    grad.set_ylabel("Applied\ngradient", fontsize=11)
    phase.plot(t, q, color="#b45335", linewidth=2.3)
    phase.text(18, 1.16, r"$q=\gamma g\delta$", ha="center", fontsize=11,
               color="#93432b")
    phase.set_ylim(-0.15, 1.55)
    phase.set_yticks([0, 1])
    phase.set_ylabel(r"$q(t)/(\gamma g\delta)$", fontsize=11)
    phase.set_xlabel("Time (ms)", fontsize=11)
    for ax in axes:
        ax.axvline(refocus, color="#a0aab1", linewidth=0.8, linestyle=":")
        ax.axvline(echo, color="#a0aab1", linewidth=0.8, linestyle=":")
        ax.set_xlim(-3, 65)
        ax.spines[["top", "right"]].set_visible(False)
    fig.text(0.5, 0.082,
             r"The 180° pulse reverses the encoding sign; the second gradient returns $q(t)$ to zero.",
             ha="center", fontsize=10, color="#45535d")
    fig.text(0.5, 0.035,
             r"For free diffusion: $b=\int q(t)^2dt=(\gamma g\delta)^2(\Delta-\delta/3)$; $S/S_0=e^{-bD}$.",
             ha="center", fontsize=10, color="#45535d")
    fig.savefig(Path.cwd() / (Path(__file__).stem + ".png"), dpi=100,
                facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()