"""(c)osanshouo 2020, CC-BY 4.0"""

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
plt.rcParams["text.usetex"] = True
plt.rcParams["font.family"] = "DejaVu Serif"
plt.rcParams["font.size"] = 14

fig = plt.figure()
plt.subplots_adjust(left=0.14, right=0.86, top=0.95, bottom=0.14)
ax = fig.add_subplot(111)

n = 0
dof = 0.04
x = np.logspace(0, 3, num=256)

def y_eq(x):
    return 45./(4.*np.sqrt(2.*np.pi**7))*dof*np.sqrt(x**3) * np.exp(-x)

def f(x, y, lam):
    return - lam/x**(n+2)*( y**2 - y_eq(x)**2 )

for lam, label, ls, zorder in [
    (1e10, r"$\lambda = 10^{10}$", "-", 9),
    (1e12, r"$\lambda = 10^{12}$", "--", 8),
    (1e14, r"$\lambda = 10^{14}$", "-.", 7),
]:
    y = [y_eq(x[0])]

    for i in range(1, len(x)):
        dx = x[i] - x[i-1]
        y.append(
            y[-1] + dx * f(x[i], y[-1], lam) / (1. + dx*2.*lam/x[i]**(n+2)*y[-1] )
        )
    y = np.array(y)

    plt.plot(x, y, zorder=zorder, ls=ls, label=label)

ax.plot(x, y_eq(x), zorder=5, ls=":", label=r"$Y_\mathrm{EQ}$")

ax.set_xlabel(r"$x = m c^2 / k_\mathrm{B} T$", fontsize=18)
ax.set_ylabel(r"$Y = n / s$", fontsize=18)
ax.set_xscale("log")
ax.set_yscale("log")
ax.set_xlim([5, 1e3])
ax.set_ylim([1e-13, 5e-3])
locmin = mpl.ticker.LogLocator(base=10.0,subs=(0.2,0.4,0.6,0.8),numticks=12)
ax.yaxis.set_minor_locator(locmin)
ax.yaxis.set_minor_formatter(mpl.ticker.NullFormatter())
ax.grid()
ax.grid(which="minor", ls=":", c="lightgray")
ax.legend(fontsize=14)

plt.savefig("cold_relic_abundance.svg")
plt.close()