from __future__ import annotations

from typing import Callable

import numpy as np


LRCallable = Callable[[int], float]


def plot_lr_curve(
    schedule: LRCallable, epochs: int, steps_per_epoch: int, title: str
) -> None:
    import matplotlib.pyplot as plt

    total_steps = int(epochs) * int(steps_per_epoch)
    steps = np.arange(total_steps)

    lrs = np.array([float(schedule(int(i))) for i in steps], dtype=np.float32)

    plt.figure(figsize=(14, 4))
    plt.plot(steps, lrs)
    plt.title(f"{title} | epochs={epochs}, steps/epoch={steps_per_epoch}")
    plt.xlabel("Global step")
    plt.ylabel("Learning rate")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()
