"""Illustrative two-state kinetics; Python 3.14, NumPy 2.3.5, Matplotlib 3.10.7.

Run from output directory; writes internal-conversion-populations.png to cwd.
"""
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


def main():
    t = np.geomspace(1e-14, 1e-7, 800)
    k21, k1 = 1e12, 1e8
    n2 = np.exp(-k21 * t)
    n1 = k21 / (k21 - k1) * np.exp(-k1 * t) * (-np.expm1(-(k21 - k1) * t))
    fig, ax = plt.subplots(figsize=(8, 3.5), dpi=120, facecolor='white')
    ax.semilogx(t * 1e12, n2, color='#225ea8', lw=2, label=r'Upper state $S_2$')
    ax.semilogx(t * 1e12, n1, color='#b44d2b', lw=2, label=r'Lower state $S_1$')
    ax.axvline(1, color='0.5', ls='--', lw=0.9)
    ax.axvline(1e4, color='0.5', ls='--', lw=0.9)
    ax.text(1.3, .55, '1 ps transfer scale', fontsize=9)
    ax.text(150, .55, '10 ns lower-state lifetime', fontsize=9)
    ax.set(xlim=(.01, 1e5), ylim=(0, 1.04), xlabel='Time after excitation (ps, logarithmic scale)',
           ylabel='Population / initial upper-state population',
           title='Fast internal conversion can precede slower excited-state decay')
    ax.legend(loc='upper left', fontsize=9)
    ax.grid(alpha=.2)
    fig.tight_layout()
    fig.savefig(Path('internal-conversion-populations.png'), facecolor='white', dpi=120)
    plt.close(fig)


if __name__ == '__main__':
    main()