"""Sequential excited-state populations; run in the PNG output directory.

Tested with Python 3.14, NumPy and Matplotlib from the root pyproject.toml.
The examples solve dn2/dt=-a*n2, dn1/dt=k21*n2-b*n1 exactly.
"""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

t = np.linspace(0, 5, 1001)
fig, axes = plt.subplots(1, 2, figsize=(8, 3.5), dpi=120, sharey=True)
fig.patch.set_facecolor('white')
for ax, a, b in zip(axes, [1., 10.], [1., 1.]):
    k21 = a
    n2 = np.exp(-a*t)
    n1 = k21*t*np.exp(-a*t) if a == b else k21/(a-b)*(np.exp(-b*t)-np.exp(-a*t))
    peak = 1/a if a == b else np.log(a/b)/(a-b)
    height = k21*peak*np.exp(-a*peak) if a == b else k21/(a-b)*(np.exp(-b*peak)-np.exp(-a*peak))
    ax.set_facecolor('white')
    ax.plot(t, n2, color='#285f9e', label='Upper state $n_2/N_0$')
    ax.plot(t, n1, color='#b65021', label='Lower state $n_1/N_0$')
    ax.plot(peak, height, 'o', color='#b65021', ms=4)
    ax.annotate(f'Peak at {peak:.3g}', (peak, height), xytext=(1.25, .74),
                arrowprops={'arrowstyle': '->', 'color': '#555555'}, fontsize=9)
    ax.set_title(rf'$a=k_{{21}}={a:g},\ b={b:g}$', fontsize=11)
    ax.set_xlabel('Time (arbitrary unit)')
    ax.set_xlim(0, 5)
    ax.set_ylim(0, 1.04)
    ax.grid(alpha=.18)
axes[0].set_ylabel('Fraction of initial population')
axes[0].legend(fontsize=8, loc='upper right')
fig.suptitle('Sequential populations: equal rates have a finite solution', fontsize=12)
fig.tight_layout(rect=(0, 0, 1, .93))
fig.savefig(Path.cwd() / 'sequential-rate-coincidence.png', facecolor='white')
plt.close(fig)