(import (srfi :43)) ; vector extensions
(load "../common.scm")

;; naive version, as per the example
; simulate lanterfish growth and returns the number of fishes after given days
;(define (simulate-growth fishes days)
;  (let days-loop [(day 0) (fishes fishes)]
;    (if (= day days)
;      (length fishes)
;      (days-loop
;	(+ day 1)
;	(let grow [(new '()) (old fishes)]
;	  (if (null? old)
;	    new
;	    (if (= (car old) 0)
;	      (grow (cons 6 (cons 8 new)) (cdr old))
;	      (grow (cons (- (car old) 1) new) (cdr old)))))))))

;; optimized version
;; instead of keeping track of fishes
;; we just keep track of how many are living at each age.
;; so, the living vector is actually indexed by age
(define (sum living)
  (vector-fold (lambda (i sum value) (+ sum value)) 0 living))

; returns a new vector with the new number of fishes at each age
(define (grow living)
  (let [(new (make-vector 9 0))]
    (vector-for-each
      (lambda (age count)
	(if (= age 0)
	  (begin
	    (vector-set! new 8 count)
	    (vector-set! new 6 count))
	  (vector-set! new (- age 1)
		       (+ count (vector-ref new (- age 1))))))
      living)
    new))

; simulate lanterfish growth and returns the number of fishes after given days
(define (simulate-growth fishes days)
  (let [(living-per-day (make-vector 9 0))]
    (let init-loop [(fishes fishes)]
      (when (not (null? fishes))
	(vector-set! living-per-day (car fishes)
		     (+ 1 (vector-ref living-per-day (car fishes))))
	(init-loop (cdr fishes))))
    (let loop [(day 0)
	       (living-per-day living-per-day)]
      (if (= day days)
	(sum living-per-day)
	(loop (+ day 1) (grow living-per-day))))))

(call-with-input-file
  "input"
  (lambda (file)
    (let [(input (read-comma-separated-numbers file))]
      (printf "part 1:~% count after 80 days: ~a~%"
	      (simulate-growth input 80))
      (printf "part 2:~% count after 256 days: ~a~%"
	      (simulate-growth input 256)))))