Advent of code 2024 day 14

import hhoppe_tools as hh
import mediapy as media

def day14(s, *, part2=False, shape=(103, 101), steps=100, visualize=False, rep=2, fps=50):
  values = np.array([re.findall(r'[-\d]+', line) for line in s.splitlines()], int)
  yxs0, dyxs = values[:, 1::-1], values[:, 3:1:-1]

  if not part2:
    yxs = (yxs0 + steps * dyxs) % shape
    count = np.zeros(shape, int)
    for yx in yxs:
      count[tuple(yx)] += 1
    ym, xm = shape[0] // 2, shape[1] // 2
    quadrants = [[g[:, :xm], g[:, xm + 1 :]] for g in [count[:ym], count[ym + 1 :]]]
    return math.prod(q.sum() for row in quadrants for q in row)

  def variance(t: int, axis: int) -> float:
    return np.var((yxs0[:, axis] + t * dyxs[:, axis]) % shape[axis])

  # We expect the pixel positions in the "tree" image to have lowest variance in both Y and X.
  moduli = shape
  remainders = [min(range(moduli[axis]), key=lambda t: variance(t, axis)) for axis in range(2)]
  t = hh.solve_modulo_congruences(remainders, moduli)

  if visualize:
    images = []
    for step in range(t + 1):
      if step > 100 and step % moduli[0] != remainders[0] and step % moduli[1] != remainders[1]:
        continue
      yxs = (yxs0 + step * dyxs) % shape
      mask = np.full(shape, False)
      mask[tuple(yxs.T)] = True
      image = hh.to_image(mask, 250, 0).repeat(rep, 0).repeat(rep, 1)
      text = f'Step{step:5}'
      hh.overlay_text(image, (10, 120), text, fontsize=14, background=250, align='tl', margin=3)
      images.append(image)
    images = [images[0]] * int(fps * 1.5) + images + [images[-1]] * int(fps * 3)
    media.show_video(images, codec='gif', fps=fps, title='day14')

  return t