import time
import numpy as np

from pymatgen.optimization.linear_assignment import LinearAssignment


def benchmark(
    sizes: list[int],
    n_repeats: int = 3,
    rng_seed: int = 42,
) -> None:
    rng = np.random.default_rng(rng_seed)

    results = []

    for n in sizes:
        times = []
        for _ in range(n_repeats):
            costs = rng.random((n, n))
            start = time.perf_counter()
            lap = LinearAssignment(costs)
            end = time.perf_counter()
            times.append(end - start)

            # sanity check
            if lap.solution is None or len(lap.solution) != n:
                raise RuntimeError(f"Invalid solution for n={n}")

        avg_time = np.mean(times)
        std_time = np.std(times)
        results.append((n, avg_time, std_time))

    print(f"{'n':>6} | {'avg_time (s)':>12} | {'std_dev (s)':>12}")
    print("-" * 36)
    for n, avg, std in results:
        print(f"{n:6d} | {avg:12.6f} | {std:12.6f}")


if __name__ == "__main__":
    benchmark(sizes=[100, 500, 1000, 5000], n_repeats=5)