(load "../common.scm")

; matrix visitor
; calls function with all coordinates of matrix and value at corresponding point
; value returned by function is set at the visiting coordinate
(define (visit-matrix! matrix function)
  (let [(data (matrix-data matrix))
	(width-1 (- (matrix-width matrix) 1))
	(height-1 (- (matrix-height matrix) 1))]
    (let y-loop [(y 0)]
      (let x-loop [(x 0)]
	(let [(value (vector-ref (vector-ref data y) x))]
	  (vector-set! (vector-ref data y) x
		       (function x y value))
	  (when (< x width-1)
	    (x-loop (+ x 1)))))
      (when (< y height-1)
	(y-loop (+ y 1))))))

; increase by 1 value at x,y in given matrix
; return that value too
(define (matrix-inc! matrix x y)
  (let [(v (+ 1 (matrix-get matrix x y 0)))]
    (matrix-set! matrix x y v)
    v))

(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
    (visit-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
    (visit-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))))