# Day 3: Spiral Memory
from time import perf_counter
from u.colors import cyan, purple, red

def d3a(limit, target):
  if target == 1:
    return 0

  for side in range(3, limit + 1, 2):
    if (ring_max := side**2) >= target:
      # Odd squares are the outer corners of each ring.
      # From there the Manhattan distance is just the ring radius plus the
      # shortest walk to one of the side midpoints.
      radius = side // 2
      side_span = side - 1
      offset = (target - (ring_max - radius)) % side_span
      return radius + min(offset, side_span - offset)

def d3b(limit, target):
  grid = {(0, 0): 1}
  neighbors = [
    (1, 0), (1, -1), (0, -1), (-1, -1),
    (-1, 0), (-1, 1), (0, 1), (1, 1),
  ]

  def cnt(x, y):
    # Only already-written cells contribute to the new value.
    c = sum(grid[(x + dx, y + dy)] for dx, dy in neighbors if (x + dx, y + dy) in grid)
    if c > target:
      return c
    grid[(x, y)] = c

  for ring in range(0, limit, 1):
    # Build one outer ring at a time around the origin.
    # Right edge. The first ring starts immediately next to the center.
    if ring == 0:
      for y in range(0, -ring - 2, -1):
        if count := cnt(ring + 1, y):
          return count
    else:
      for y in range(ring, -ring - 2, -1):
        if count := cnt(ring + 1, y):
          return count

    # Bottom edge.
    for x in range(ring + 1, -ring - 2, -1):
      if count := cnt(x, -ring - 1):
        return count

    # Left edge.
    for y in range(-ring - 1, ring + 2, 1):
      if count := cnt(-ring - 1, y):
        return count

    # Top edge.
    for x in range(-ring - 1, ring + 2, 1):
      if count := cnt(x, ring + 1):
        return count

t = perf_counter()
cyan(f'\na) {d3a(1000, 289326)}')
purple(f'b) {d3b(1000, 289326)}')
red(f'\nT: {(perf_counter() - t) * 1000} ms\n')