#lang curly-fn racket

(require "scripts/aoc.rkt")

;; solution for day 14

(define (part1 input)
  (define (parse-mask msk)
    (define orval (for/fold ([val 0]) ([chr (in-string msk)])
                    (+ (* val 2) (if (equal? chr #\1) 1 0))))
    (define andval (for/fold ([val 0]) ([chr (in-string msk)])
                     (+ (* val 2) (if (equal? chr #\0) 0 1))))
    (values orval andval))
  (define orval 0)
  (define andval (sub1 (expt 2 37)))
  (define mem (make-hash))
  (for ([ins (in-list input)])
    (match ins
      [(cons 'mask mask)
       (define-values [ov av] (parse-mask mask))
       (set! orval ov) (set! andval av)]
      [(cons 'set (cons addr val))
       (define new-val (bitwise-ior (bitwise-and val andval) orval))
       (hash-set! mem addr new-val)]))
  (for/sum ([(k v) (in-hash mem)])
    v))

(define (addr->bits addr)
  (vector-copy 
  (for/vector ([i (in-range 36)])
    (bitwise-and 1 (arithmetic-shift addr (+ -35 i))))))
(define (bits->addr bits)
  (for/fold ([val 0]) ([bit (in-vector bits)])
    (+ (* val 2) bit)))

(define (part2 input)
  (define mem (make-hash))
  (define mask #f)
  (for ([ins (in-list input)])
    (match ins
      [(cons 'mask m) (set! mask m)]
      [(cons 'set (cons addr val))
       (define addr-bits (addr->bits addr))
       (for ([chr (in-string mask)] [i (in-naturals)])
         (when (equal? chr #\1)
           (vector-set! addr-bits i 1)))
       (define addrs-set (list addr-bits))
       (for ([chr (in-string mask)] [i (in-naturals)])
         (when (equal? chr #\X)
           (define ones (map (lambda (x) (define y (vector-copy x)) (vector-set! y i 1) y)
                             addrs-set))
           (define zeros (map (lambda (x) (define y (vector-copy x)) (vector-set! y i 0) y)
                             addrs-set))
           (set! addrs-set (append ones zeros))))
       (for ([addr (in-list addrs-set)])
         (hash-set! mem (bits->addr addr) val))]))

  (for/sum ([(k v) (in-hash mem)])
    v))

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

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

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