#lang racket
(require "lib/common.rkt"
         fancy-app
         math/array
         math/matrix)

(define transition-matrix
  (matrix [[0 1 0 0 0 0 0 0 0]
           [0 0 1 0 0 0 0 0 0]
           [0 0 0 1 0 0 0 0 0]
           [0 0 0 0 1 0 0 0 0]
           [0 0 0 0 0 1 0 0 0]
           [0 0 0 0 0 0 1 0 0]
           [1 0 0 0 0 0 0 1 0]
           [0 0 0 0 0 0 0 0 1]
           [1 0 0 0 0 0 0 0 0]]))

(define (generate-initial-vector fish)
  (define vec (make-vector 9 0))
  (for ([f (in-list fish)])
    (vector-set! vec f (add1 (vector-ref vec f))))
  vec)

(define (simulate fish days)
  (define V
    (matrix* (matrix-expt transition-matrix days)
             (->col-matrix (generate-initial-vector fish))))
  (for/sum ([v (in-array V)]) v))

(define (day6a fish) (simulate fish 80))
(define (day6b fish) (simulate fish 256))

(module+ main
  (call-with-input-file "data/day6.txt"
    (λ (prt)
      (define fish
        (~> (port->string prt)
            string-split
            first
            (string-split ",")
            (map string->number _)))
      (answer 6 1 (day6a fish))
      (answer 6 2 (day6b fish)))))

(module+ test
  (require rackunit)
  (define test-fish '(3 4 3 1 2))

  (check-equal? (day6a test-fish) 5934)
  (check-equal? (day6b test-fish) 26984457539))