"""Draw original charged-vector diagrams; tested with Python 3.14, numpy, matplotlib."""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def straight(ax, start, end, label=None, offset=(0, 0.12)):
    ax.plot([start[0], end[0]], [start[1], end[1]], color="#225a8e", lw=2.2)
    a, b = np.array(start), np.array(end)
    ax.annotate("", xy=a + .62*(b-a), xytext=a + .40*(b-a),
                arrowprops={"arrowstyle": "-|>", "color": "#225a8e", "lw": 1.5})
    if label:
        p = (a+b)/2 + np.array(offset)
        ax.text(*p, label, ha="center", va="center", fontsize=12)


def photon(ax, start, end, label=None):
    a, b = np.array(start), np.array(end)
    t = np.linspace(0, 1, 300)
    tangent = b-a
    normal = np.array([-tangent[1], tangent[0]])/np.linalg.norm(tangent)
    xy = a[:, None] + tangent[:, None]*t + .045*normal[:, None]*np.sin(12*np.pi*t)
    ax.plot(xy[0], xy[1], color="#b46b00", lw=1.7)
    if label:
        p=(a+b)/2
        ax.text(p[0]+.11, p[1]+.13, label, fontsize=12, ha="center")


def dot(ax, p):
    ax.plot(*p, "o", color="black", ms=5)


def main():
    fig, axes=plt.subplots(2,2,figsize=(10,6),layout="constrained",facecolor="white")
    for ax in axes.flat:
        ax.set(xlim=(-1.3,1.3),ylim=(-1.05,1.05),aspect="equal")
        ax.axis("off")
    ax=axes[0,0]
    ax.set_title(r"Cubic vertex: $ABB^*$, order $e$",fontsize=13)
    v=(0,0)
    straight(ax,(-1,-.65),v,r"$B$")
    straight(ax,v,(1,-.65),r"$B$")
    photon(ax,v,(0,.8),r"$A$")
    dot(ax,v)
    ax=axes[0,1]
    ax.set_title(r"Quartic vertex: $AABB^*$, order $e^2$",fontsize=13)
    straight(ax,(-1,-.65),v,r"$B$")
    straight(ax,v,(1,-.65),r"$B$")
    photon(ax,v,(-.8,.75),r"$A$")
    photon(ax,v,(.8,.75),r"$A$")
    dot(ax,v)
    for ax, crossed, channel in [(axes[1,0],False,"t"),(axes[1,1],True,"u")]:
        ax.set_title(rf"Identical $B^+B^+$: ${channel}$ channel, order $e^2$",fontsize=13)
        top,bottom=(0,.45),(0,-.45)
        straight(ax,(-1,.85),top,r"$p_1$",(-.05,.13))
        straight(ax,(-1,-.85),bottom,r"$p_2$",(-.05,-.15))
        straight(ax,top,(1,.85),r"$p_4$" if crossed else r"$p_3$",(.05,.13))
        straight(ax,bottom,(1,-.85),r"$p_3$" if crossed else r"$p_4$",(.05,-.15))
        photon(ax,bottom,top,r"$\gamma$")
        dot(ax,top);dot(ax,bottom)
    fig.savefig(Path.cwd()/Path(__file__).with_suffix(".png").name,dpi=120,facecolor="white",transparent=False)
    plt.close(fig)


if __name__=="__main__":
    main()