#lang curly-fn racket

(require "scripts/aoc.rkt")

;; solution for day 17

(struct pt [x y z w] #:transparent)
(define pt-getters (list pt-x pt-y pt-z pt-w))
(define origin (pt 0 0 0 0))

(define (pt+ . args)
  (for/fold ([ans (pt 0 0 0 0)]) ([x (in-list args)])
    (apply pt (map #{+ (% ans) (% x)} pt-getters))))

(define (in-pts lower upper use-w?)
  (define (in-range+ getter)
    (in-range (sub1 (getter lower)) (+ 2 (getter upper))))
  (for*/stream ([x (in-range+ pt-x)]
                [y (in-range+ pt-y)]
                [z (in-range+ pt-z)]
                [w (if use-w? (in-range+ pt-w) (in-range 1))])
    (pt x y z w)))

(define nearby/4d
  (for/list ([pt (in-pts origin origin #t)] #:unless (equal? pt origin)) pt))
(define nearby/3d (filter #{zero? (pt-w %)} nearby/4d))

(define (calc-bounds world)
  (define (helper getter reducer)
    (apply reducer (set-map world getter)))
  (define (helper2 reducer)
    (apply pt (map #{helper % reducer} pt-getters)))
  (values (helper2 min) (helper2 max)))

(define (simulate world use-w?)
  (define new-world (mutable-set))
  (define-values [lower upper] (calc-bounds world))
  (define nearby (if use-w? nearby/4d nearby/3d))

  (for ([pt (in-pts lower upper use-w?)])
    (define ct
      (for*/sum ([nb (in-list nearby)])
        (if (set-member? world (pt+ pt nb)) 1 0)))
    (when (if (set-member? world pt) (<= 2 ct 3) (= 3 ct))
      (set-add! new-world pt)))

  new-world)

(define (simulate/times world times use-w?)
  (set-count
    (for/fold ([world world]) ([_ (in-range times)])
      (simulate world use-w?))))

(define part1 #{simulate/times % 6 #f})
(define part2 #{simulate/times % 6 #t})

;; parse input file
(define (parse fname)
  (define input (file->lines fname))
  (define world (mutable-set))
  (for ([line (in-list input)] [y (in-naturals)])
    (for ([char (in-string line)] [x (in-naturals)])
      (when (char=? #\# char)
        (set-add! world (pt x y 0 0)))))
  world)

(module+ test
  (require rackunit)
  ;; tests here
  (displayln "no tests :("))

(module+ main
  (define input (parse "inputs/17"))
  (answer 17 1 (time (part1 input)))
  (answer 17 2 (time (part2 input)))
  (displayln "meow"))