(library (lib)
  (export solve-part1 solve-part2)
  (import (chezscheme))

  ; Is this direction up or down?
  (define (vertical? dir)
    (= 0 (vector-ref dir 0)))

  ; Rotate stepping direction counterclockwise 90 degrees
  (define (rotate-ccw dir)
    (cond
      ((equal? dir '#(1 0))
       '#(0 1))
      ((equal? dir '#(0 1))
       '#(-1 0))
      ((equal? dir '#(-1 0))
       '#(0 -1))
      ((equal? dir '#(0 -1))
       '#(1 0))))

  ; Follow the spiral to find (x,y) `pos' of the `address'th square.
  ; The idea is to take `steps' steps in the current direction `dir',
  ; rotate `dir', take `steps' steps again, rotate, increment `steps',
  ; then start over, and so on. So the sequence starts like this:
  ;   1 step right, 1 step up,
  ;   2 steps left, 2 steps down,
  ;   3 steps right, 3 steps up,
  ;   4 steps left, 4 steps down,
  ;   etc.
  ; Note: for each value of `steps', the last stage is always up/down.
  (define (get-position address)
    (let loop ((pos '#(0 0))
               (dir '#(1 0))
               (steps 1)
               (count 0)
               (total 1))
      (if (= total address)
          pos
          (if (= count (- steps 1))
              ; Yes, this is the last step before we need to turn
              (loop
                ; Take the step, i.e. add direction to position
                (vector-map + pos dir)
                ; Rotate counterclockwise after this step
                (rotate-ccw dir)
                ; If this step is vertical, we need to increment
                (+ steps (if (vertical? dir) 1 0))
                ; Reset step counter for the current direction
                0
                ; Keep track of total steps since square one
                (+ total 1))
              ; No, this is an ordinary step, no rotation needed
              (loop
                (vector-map + pos dir)
                dir
                steps
                (+ count 1)
                (+ total 1))))))

  ; Given a position, return all eight positions adjacent to it
  (define (get-adjacent pos)
    (map
      (lambda (dir) (vector-map + pos dir))
      '(#(1 0) #(1 1) #(0 1) #(-1 1) #(-1 0) #(-1 -1) #(0 -1) #(1 -1))))

  ; Given a position, sum the values of all initialized adjacent squares
  (define (sum-adjacent pos memory)
    (fold-left
      (lambda (sum adj)
        ; Try to retrieve `pair' from memory
        (let ((pair (assoc adj memory)))
          (+ sum
             ; If position `(car pair)' has been written,
             ; add its value `(cdr pair) to the sum total.
             (if pair (cdr pair) 0))))
      0
      (get-adjacent pos)))

  (define (solve-part1 target)
    (let* ((xy (get-position target))
           (x (vector-ref xy 0))
           (y (vector-ref xy 1)))
      (+ (abs x) (abs y))))

  (define (solve-part2 target)
    (let loop ((address 2)
               (memory '((#(0 0) . 1))))
      (let ((value (cdar memory)))
        ; Has the most recently written `value' crossed the threshold?
        (if (> value target)
            value
            ; If not, move on, writing a new `val' to the next `pos'
            (let* ((pos (get-position address))
                   (val (sum-adjacent pos memory)))
              (loop
                (+ address 1)
                (cons (cons pos val) memory)))))))

)