(load "../common.scm")

(define-record-type vec2i (fields x y))

; flash given octopus
; ie. increase energy of neighbor ones
; returns list of octopuses that now need to flash too
(define (flash! matrix x y)
  (let [(extra '())]
    (let y-loop [(yy (- y 1))]
      (let x-loop [(xx (- x 1))]
	(when (not (and (= xx x) (= yy y)))
	  (when (< (matrix-get matrix xx yy 0) 10)
	    (when (= (matrix-inc! matrix xx yy) 10)
	      (set! extra (cons (make-vec2i xx yy) extra)))))
	(when (< xx (+ x 1))
	  (x-loop (+ xx 1))))
      (when (< yy (+ y 1))
	(y-loop (+ yy 1))))
    extra))

(define (step! matrix)
  (let [(flashes '())]
    ; increase energy of all octopuses
    (update-matrix! 
      matrix
      (lambda (x y energy)
	(when (= energy 9)
	  (set! flashes (cons (make-vec2i x y) flashes)))
	(+ 1 energy)))
    ; recursively flash octopuses
    (let loop [(flash flashes)]
      (when (not (null? flash))
	(loop (append (cdr flash)
		      (flash! matrix (vec2i-x (car flash))
			      (vec2i-y (car flash)))))))
    ; drain all octopuses that flashed
    (update-matrix! 
      matrix
      (lambda (x y energy)
	(if (< energy 10) energy 0)))
    ; count octopuses that flashed
    (vector-fold
      (lambda (y sum line)
	(+ sum
	   (vector-fold
	     (lambda (x sum value)
	       (if (= value 0) (+ sum 1) sum))
	     0 line)))
      0 (matrix-data matrix))))

(call-with-input-file
  "input"
  (lambda (file)
    (let [(matrix (load-matrix file))
	  (flash-count 0)]
      (let loop [(i 0)]
	(set! flash-count (+ flash-count
			     (step! matrix)))
	(when (< i 99)
	  (loop (+ i 1))))
      (printf "part 1:~% ~a flashes after a hundred steps.~%"
	      flash-count)
      (let loop [(i 100)]
	(if (= (matrix-sum matrix) 0)
	  (printf "part 2:~% Octopuses synchronized at step ~a.~%" i)
	  (begin
	    (step! matrix)
	    (loop (+ i 1))))))))