#lang curly-fn racket

(require "scripts/aoc.rkt")
;; data/bit-vector is hella nerfed so.... rosette uwu
;; rosette is slow af for this but whatever
(require (only-in rosette bv bvand bvor bvxor rotate-left concat bitvector zero-extend
                  bitvector->natural))

;; solution for day 14
(define (mask->bv bits select)
  (apply concat (map #{if (equal? % select) (bv select 1) (bv (- 1 select) 1)} bits)))

(define (pos->bv pos)
  (rotate-left (- 35 pos) (zero-extend (bv 1 1) (bitvector 36))))

(define (part1 input)
  (define orval #f)
  (define andval #f)
  (define mem (make-hash))
  (for ([ins (in-list input)])
    (match ins
      [(list 'mask bits ...)
       (set! orval (mask->bv bits 1))
       (set! andval (mask->bv bits 0))]
      [(list 'set addr val)
       (hash-set! mem addr
                  (bitvector->natural (bvor (bvand (bv val 36) andval) orval)))]))
  (for/sum ([(k v) (in-hash mem)])
    v))

(define (part2 input)
  (define mem (make-hash))
  (define mask #f)
  (for ([ins (in-list input)])
    (match ins
      [(list 'mask bits ...) (set! mask bits)]
      [(list 'set addr val)
       (define addr-bits (bvor (bv addr 36) (mask->bv mask 1)))
       (define addrs-set (list addr-bits))
       (for ([item (in-list mask)] [i (in-naturals)] #:when (equal? 'x item))
         (set! addrs-set (append addrs-set (map #{bvxor (pos->bv i) %} addrs-set))))
       (for ([addr (in-list addrs-set)])
         (hash-set! mem (bitvector->natural 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 (for/list ([chr (in-string mask)])
                     (match chr
                       [#\1 1]
                       [#\0 0]
                       [#\X 'x])))]
      [(pregexp #px"^mem\\[([0-9]+)\\] = ([0-9]+)$" (list _ addr val))
       (list 'set (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"))