"""Stationary reflecting diffusion in [0,L]; original model, not measured data.

Tested with Python 3.14, NumPy and Matplotlib (repository pyproject.toml).
Write the PNG to the current directory, normally the mirrored _media directory.
For u=D*t/L**2, the stationary mean squared displacement is
16/pi**4 * sum_{n odd} (1-exp(-n*n*pi*pi*u))/n**4, in units of L**2.
This follows by expanding x-L/2 in Neumann cosine modes. Its limit is 1/6.
"""
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


def stationary_msd(u, terms=1024):
    u = np.asarray(u)
    odd = np.arange(1, 2*terms, 2, dtype=float)
    return 16/np.pi**4 * np.sum(-np.expm1(-np.pi**2*u[..., None]*odd**2)/odd**4, axis=-1)


def main():
    u = np.geomspace(1e-4, 1, 500)
    restricted = stationary_msd(u)
    assert np.all(np.diff(restricted) > 0)
    assert np.all(restricted <= 2*u)
    assert abs(float(stationary_msd(5))-1/6) < 1e-10
    assert abs(float(stationary_msd(1e-5))/(2e-5)-1) < 0.01
    fig, ax = plt.subplots(figsize=(8, 5), dpi=100, facecolor='white')
    fig.subplots_adjust(left=0.12, right=0.97, bottom=0.24, top=0.82)
    fig.suptitle('Walls limit displacement without stopping diffusion', fontsize=14, weight='bold', y=0.96)
    fig.text(0.5, 0.88, 'One-dimensional interval of width L; initially uniform equilibrium', ha='center', fontsize=11)
    ax.loglog(u, 2*u, '--', color='#b45335', label='Free diffusion: 2u', linewidth=2)
    ax.loglog(u, restricted, color='#216c9b', label='Reflecting walls', linewidth=2.6)
    ax.axhline(1/6, color='#66717b', linestyle=':', label='Confined long-time limit: 1/6')
    ax.set_xlabel(r'Time $u=D\tau/L^2$', fontsize=11)
    ax.set_ylabel(r'Mean squared displacement / $L^2$', fontsize=11)
    ax.set_xlim(1e-4, 1)
    ax.grid(which='major', color='#e1e6ea', linewidth=0.7)
    ax.legend(loc='upper left', fontsize=10)
    fig.text(0.5, 0.045, 'Original continuum calculation with perfectly reflecting walls.\n'
             'The plateau depends on equilibrium sampling of the starting position.', ha='center', fontsize=10)
    fig.savefig(Path.cwd()/(Path(__file__).stem+'.png'), dpi=100, facecolor='white', transparent=False)
    plt.close(fig)


if __name__ == '__main__':
    main()