r"""
Szwarc's living polymerization: the narrow Poisson distribution
==================================================================

Szwarc (1956) showed that anionic polymerization of styrene initiated by
sodium naphthalenide has no termination: every chain starts at once and
keeps growing while monomer lasts. Each chain then receives monomers as
independent random events, giving a Poisson distribution
(:func:`~chemistrykit.polymer.systems.living_polymerization.poisson_number_fraction`)
with :math:`\text{PDI}=1+\nu/(1+\nu)^2`
(:func:`~chemistrykit.polymer.systems.living_polymerization.poisson_pdi`) -- close to 1,
far narrower than the Flory-Schulz distribution of a step-growth
polymer with the same :math:`\bar X_n`. A seeded stochastic simulation
(:func:`~chemistrykit.polymer.systems.living_polymerization.simulate_living_polymerization`)
reproduces the prediction.
"""

# %%
import matplotlib.pyplot as plt
import numpy as np

from chemistrykit.polymer.systems.living_polymerization import (
    poisson_number_average_DP,
    poisson_number_fraction,
    poisson_pdi,
    simulate_living_polymerization,
)
from chemistrykit.polymer.systems.molecular_weight_distribution import (
    flory_schulz_number_fraction,
    flory_schulz_pdi,
)

n_chains, n_monomers = 5000, 250000
nu = n_monomers / n_chains
x_sim = simulate_living_polymerization(n_chains, n_monomers, rng=0)
pdi_sim = np.mean(x_sim**2) / np.mean(x_sim) ** 2
print(f"nu = {nu:.0f}: Xn simulated = {x_sim.mean():.2f}, Poisson = {poisson_number_average_DP(nu):.2f}")
print(f"PDI simulated = {pdi_sim:.4f}, Poisson = {poisson_pdi(nu):.4f}")

# Step-growth polymer with the same Xn = 1/(1-p)
p = 1 - 1 / poisson_number_average_DP(nu)
print(f"Flory-Schulz PDI at the same Xn: {flory_schulz_pdi(p):.3f}")

# %%
x = np.arange(1, 160)
fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
axes[0].hist(x_sim, bins=np.arange(0.5, 160.5, 2), density=True, alpha=0.5, label="stochastic simulation")
axes[0].plot(x, poisson_number_fraction(x, nu), "k-", label="Poisson (living)")
axes[0].plot(x, flory_schulz_number_fraction(x, p), "r--", label="Flory-Schulz, same Xn")
axes[0].set_xlabel("chain length x")
axes[0].set_ylabel("number fraction")
axes[0].set_title(f"Living chains, nu = {nu:.0f}")
axes[0].legend()

nus = np.logspace(0, 3, 100)
axes[1].semilogx(nus, poisson_pdi(nus), label="living (Poisson)")
axes[1].axhline(2.0, color="r", ls="--", label="step growth limit (2)")
axes[1].set_xlabel(r"$\nu$ (monomers per initiator)")
axes[1].set_ylabel("PDI")
axes[1].set_title("Polydispersity")
axes[1].legend()
plt.tight_layout()
plt.show()