#lang curly-fn racket

(require "scripts/aoc.rkt"
         parser-tools/lex
         (prefix-in : parser-tools/lex-sre)
         parser-tools/yacc)

;; solution for day 11

(define ((part div3?) input)
  (define (m-op n) (second (hash-ref input n)))
  (define (m-test n) (third (hash-ref input n)))
  (define (m-if-t n) (fourth (hash-ref input n)))
  (define (m-if-f n) (fifth (hash-ref input n)))

  (define (m-op* n worry)
    (define op (m-op n))
    ((match (first op) ['+ +] ['* *])
     (match (second op) ['old worry] [x x])
     (match (third op) ['old worry] [x x])))

  (define inspect-count (make-hash))
  (for ([k (in-hash-keys input)])
    (hash-set! inspect-count k 0))

  (define all-monkeys-modulo
    (apply lcm (for/list ([mky (in-hash-keys input)]) (m-test mky))))

  (for*/fold ([monkeys (for/hash ([(k v) (in-hash input)]) (values k (first v)))])
             ([rnd (in-range (if div3? 20 10000))] [mky (in-hash-keys input)])
    (define items (hash-ref monkeys mky))

    (for/fold ([monkeys (hash-set monkeys mky '())])
              ([item (in-list items)])
      (hash-update! inspect-count mky add1)
      (define new-worry (~> item (m-op* mky _) (#{if div3? (quotient % 3) %})
                            (remainder _ all-monkeys-modulo)))
      (define next-mky
        (if (zero? (remainder new-worry (m-test mky))) (m-if-t mky) (m-if-f mky)))
      (hash-update monkeys next-mky #{append % (list new-worry)})))

  (match-define (list* top0 top1 rst) (sort (hash-values inspect-count) >))
  (* top0 top1))

(define (parse fname)
  (define-tokens monkey-tokens (NUM))
  (define-empty-tokens monkey-tokens-empty (NEXT START OP OLD PLUS TIMES TEST IF-TRUE IF-FALSE EOF))

  (define monkey-lex
    (lexer
      [(eof) (token-EOF)]
      ["Monkey" (token-NEXT)]
      ["Starting items" (token-START)]
      ["Operation: new" (token-OP)]
      ["old" (token-OLD)]
      ["+" (token-PLUS)]
      ["*" (token-TIMES)]
      ["Test: divisible by" (token-TEST)]
      ["If true: throw to monkey" (token-IF-TRUE)]
      ["If false: throw to monkey" (token-IF-FALSE)]
      [(:+ numeric) (token-NUM (string->number lexeme))]
      [(:or whitespace ":" "=" ",") (monkey-lex input-port)]))

  (define monkey-parse
    (parser
      [start exps] [end EOF] [error void]
      [tokens monkey-tokens monkey-tokens-empty]
      [grammar
        (exps [(exp) (list $1)]
              [(exp exps) (cons $1 $2)])
        (exp [(NEXT NUM START num-list OP op-expr TEST NUM IF-TRUE NUM IF-FALSE NUM)
              (list $2 $4 $6 $8 $10 $12)])
        (num-list [(NUM) (list $1)]
                  [(NUM num-list) (cons $1 $2)])
        (op-expr [(opnd op opnd) (list $2 $1 $3)])
        (opnd [(OLD) 'old]
              [(NUM) $1])
        (op [(PLUS) '+]
            [(TIMES) '*])]))

  (define file (open-input-file fname))
  (for/hash ([item (in-list (monkey-parse (thunk (monkey-lex file))))])
    (values (first item) (rest item))))

(module+ main
  (define input (parse "inputs/11"))
  (answer 11 1 (time ((part #t) input)))
  (answer 11 2 (time ((part #f) input)))
  (displayln "meow"))