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

  ; Split list at first delimiter into `prefix' and `suffix'
  ; Return value is a pair like `((p r e f i x) s u f f i x)'
  (define (list-split-left delimiter? xs)
    (let loop ((prefix '())
               (suffix xs))
      (if (null? suffix)
        (cons prefix suffix)
        (let ((x (car suffix)))
          (if (delimiter? x)
              ; Found first delimiter, so return immediately
              (cons prefix (cdr suffix))
              ; `x' isn't a delimiter, so append it to `prefix'
              (loop (append prefix (list x)) (cdr suffix)))))))

  ; Split list at given delimiter into list of sublists
  (define (list-split delimiter? xs)
    (let loop ((pieces '())
               (rest xs))
      (if (null? rest)
          ; Fix order and remove all empty sublists from output
          ; (which are caused by consecutive delimiters in `xs')
          (reverse (remp null? pieces))
          ; Extract next piece from `rest' and prepend it to `pieces'
          (let ((next (list-split-left delimiter? rest)))
            (loop (cons (car next) pieces) (cdr next))))))

  ; Split string at whitespace into list of words
  (define (string-split str)
    (map list->string (list-split char-whitespace? (string->list str))))

  ; Find the index at which `needle' occurs in `haystack', or `#f'
  (define (find-index needle haystack)
    (let loop ((xs haystack) (i 0))
      (if (null? xs)
          #f
          (if (equal? needle (car xs))
              i
              (loop (cdr xs) (+ i 1))))))

  ; `apply' but for vectors, assuming `proc' returns a scalar
  (define (vector-apply proc vec)
    (apply proc (vector->list vec)))

  ; Redistribute the items according to the puzzle description
  (define (redistribute old)
    (let* ((n-max (vector-apply max old))                ; largest value in `old'
           (i-max (find-index n-max (vector->list old))) ; index of `n-max' in `old'
           (new (vector-apply vector old))               ; allocate copy of `old'
           (len (vector-length old)))                    ; number of slots in `old'
      ; Clear the slot that had the most items
      (vector-set! new i-max 0)
      ; Add those items to the other slots one by one
      (let loop ((i (mod (+ i-max 1) len))
                 (remaining n-max))
        (if (= remaining 0)
            new
            (let ((n (vector-ref new i)))
              (vector-set! new i (+ n 1))
              (loop (mod (+ i 1) len) (- remaining 1)))))))

  ; Keep redistributing items until a previous state is revisited
  (define (solve-puzzle str)
    (let ((init (list->vector (map string->number (string-split str)))))
      (let loop ((states (list init)))
        (if (find-index (car states) (cdr states))
            states
            (loop (cons (redistribute (car states)) states))))))

  ; Loop detected after `(length states)' cycles, minus initial state
  (define (solve-part1 str)
    (- (length (solve-puzzle str)) 1))

  ; Loop start state first seen somewhere in `(cdr states)', at which index?
  (define (solve-part2 str)
    (let ((states (solve-puzzle str)))
      (+ (find-index (car states) (cdr states)) 1)))

)