"""Original SGDR schedule example; write the PNG in the current directory.

Tested with Python 3.14.4, NumPy 2.3.5 and Matplotlib 3.10.7.
Only a declared learning-rate schedule is plotted, not training observations.
"""

from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D


eta_min, eta_max = 0.0, 0.1
durations = [4, 8, 16]
starts = np.concatenate(([0], np.cumsum(durations)[:-1]))
assert list(starts) == [0, 4, 12]


def rate(elapsed, duration):
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + np.cos(np.pi * elapsed / duration))


assert np.isclose(rate(0, 4), 0.1)
assert np.isclose(rate(2, 4), 0.05)
assert np.isclose(rate(4, 4), 0.0)
assert np.isclose(rate(1, 4), 0.08535533905932738)

plt.rcParams.update({"font.size": 11, "figure.facecolor": "white"})
fig, ax = plt.subplots(figsize=(9.0, 4.0), dpi=100)
fig.subplots_adjust(left=0.09, right=0.98, top=0.79, bottom=0.21)
fig.suptitle("Cosine cycles with growing warm-restart periods", fontsize=15, y=0.96)
colors = ["#1766a0", "#ba5518", "#237b43"]
for index, (begin, duration, color) in enumerate(zip(starts, durations, colors, strict=True)):
    elapsed = np.linspace(0, duration, 401)
    ax.plot(begin + elapsed, rate(elapsed, duration), color=color, lw=2.4)
    ax.scatter([begin], [eta_max], color=color, s=34, zorder=5)
    ax.scatter([begin + duration], [eta_min], facecolors="white", edgecolors=color,
               s=44, zorder=5)
    ax.text(begin + duration / 2, 0.113, f"Cycle {index + 1}: {duration} epochs",
            ha="center", fontsize=10)
    if begin:
        ax.vlines(begin, eta_min, eta_max, color="#999999", ls=":", lw=1.0)
ax.scatter([28], [eta_max], color="#555555", s=34, zorder=5)
ax.vlines(28, eta_min, eta_max, color="#999999", ls=":", lw=1.0)
ax.set(xlim=(-0.4, 28.5), ylim=(-0.005, 0.125),
       xlabel="Elapsed training epochs (continuous schedule)", ylabel="Learning rate η")
ax.set_xticks([0, 4, 8, 12, 16, 20, 24, 28])
ax.set_yticks([0, 0.025, 0.05, 0.075, 0.1])
ax.grid(alpha=0.16)
ax.spines[["top", "right"]].set_visible(False)
ax.legend(handles=[
    Line2D([], [], marker="o", color="none", markerfacecolor="white", markeredgecolor="#555555",
           label="End-of-cycle limit"),
    Line2D([], [], marker="o", color="none", markerfacecolor="#555555", markeredgecolor="#555555",
           label="Rate at restart"),
], loc="upper right", frameon=False, fontsize=9)
fig.text(0.5, 0.055, "Chosen schedule: bounds 0 and 0.1; durations 4, 8, 16; current parameters retained.",
         ha="center", fontsize=10)
fig.savefig(Path("sgdr-cosine-restart-cycles.png"), facecolor="white")
plt.close(fig)