#lang curly-fn racket

(require "scripts/aoc.rkt")

;; solution for day 14

(define (part1 input)
  (define orval #f)
  (define andval #f)
  (define mem (make-hash))
  (for ([ins (in-list input)])
    (match ins
      [(list 'mask mask)
       (set! orval (string->number (string-replace mask "X" "0") 2))
       (set! andval (string->number (string-replace mask "X" "1") 2))]
      [(list 'set addr val)
       (hash-set! mem addr (bor (band val andval) orval))]))
  (apply + (hash-values mem)))

(define (part2 input)
  (define mem (make-hash))
  (define orval #f)
  (define xorvals #f)
  (for ([ins (in-list input)])
    (match ins
      [(list 'mask mask)
       (set! orval (string->number (string-replace mask "X" "0") 2))
       (set! xorvals
         (for/fold ([xorvals '(0)]) ([i (in-naturals)] [m (in-string mask)] #:when (char=? #\X m))
          (append xorvals (map #{bor (bshl 1 (- 35 i)) %} xorvals))))]
      [(list 'set addr val)
       (for ([xv (in-list xorvals)])
         (hash-set! mem (bxor (bor orval addr) xv) val))]))
  (apply + (hash-values mem)))

;; parse input file
(define (parse fname)
  (define input (file->lines fname))
  (for/list ([i (in-list input)])
    (match i
      [(pregexp #px"^mask = (.*?)$" (list _ mask))
       (list 'mask mask)]
      [(pregexp #px"^mem\\[([0-9]+)\\] = ([0-9]+)$" (list _ addr val))
       (list 'set (string->number addr) (string->number val))]
      [_ (error "not shonks i")])))

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