test_input <- c(
  "mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X",
  "mem[8] = 11",
  "mem[7] = 101",
  "mem[8] = 0"
)

real_input <- readLines('input.txt')

library(stringr)  # For str_split(), str_extract()


# Given a number `n`, returns a numeric representation of the 36 consituent bits
to_bits <- function(n) {
  bits <- as.numeric(intToBits(n))
  c(rep(0, 4), rev(bits))
}


# Given a binary vector `bin_vec` as produced by `to_bits()`, return the
# decimal value
to_numeric <- function(bin_vec) {
  places <- (length(bin_vec):1) - 1
  sum(2^places[which(bin_vec == 1)])
}


# Given a string `s` in the form of 'mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X',
# return the part to the right of the `=` sign as a character vector
to_mask <- function(s) {
  str_split(str_extract(s, '\\b[X\\d]+$'), '', simplify = T)
}


# Given a binary vector `bits`, the bits vector after applying a mask 
# `masked_bits`, and the numeric representation of the masked bits `result`,
# returns a named list containing those items
output_el <- function(bits, masked_bits, result) {
  list(bits = bits, masked_bits = masked_bits, result = result)
}

# Given a character vector representing a bit mask `mask` and a vector of 
# bits `vec`, applies the mask to the bit vector as described in the puzzle
# instructions
apply_mask <- function(mask, vec) {
  vec[which(mask != 'X')] <- as.numeric(mask[mask != 'X'])
  vec
}


input <- real_input   # Parse the input
mask <- rep('X', 36)  # A dummy mask
output <- list()      # List to hold the output

# For each line in the input lines
for (line in input) {
  type <- str_extract(line, '^\\w+')  # The 'type' is the word to the left of `=`
  if (type == 'mask') { mask <- to_mask(line)}
  if (type == 'mem') {
    # For 'mem' lines, convert the number to 'bits', apply the mask, convert
    # back to a number, then save the result to the `output` list  with a list
    # item name from the memory address
    address <- str_extract(line, '(?<=\\[)\\d+(?=\\])')
    bits <- to_bits(as.numeric(str_extract(line, '\\d+$')))
    masked_bits <- apply_mask(mask, bits)
    result <- to_numeric(masked_bits)
    output[[address]] <- output_el(bits, masked_bits, result)
  }
}
results <- sapply(output, function(x){ x$result })  # Get all the numeric results
answer1 <- sum(results)

source('exercise_1.R')

test_input2 <- c(
  "mask = 000000000000000000000000000000X1001X",
  "mem[42] = 100",
  "mask = 00000000000000000000000000000000X0XX",
  "mem[26] = 1"
)


# Given a string `s` in the format "mask = 000000000000000000000000000000X1001X",
# return a character matrix where each row represents a mask to apply.
to_mask <- function(s) {
  mask_chars <- unlist(str_split(str_extract(s, '\\b[X\\d]+$'), ''))
  matrix(unlist(all_masks(mask_chars)), ncol = 36, byrow = T)
}


# Given a character vector consisting of 1's, 0's and X's `chars` (the `i`
# argument should be ignored, it's used for the recursion), return a list
# containing all possible interpretations of the mask
all_masks <- function(chars, i = 1) {
  if (i > length(chars)){ return(chars) }
  if (chars[i] == 'X') {
    chars[i] <- 'Z'; chars0 <- chars
    chars[i] <- '1'; chars1 <- chars
    list(all_masks(chars0, i + 1), all_masks(chars1, i + 1))
  } else {
    all_masks(chars, i + 1)
  }
}


# Given a character vector representing a bit mask to apply `mask` and a 
# 'binary' vector `vec`, return a binary vector representing the result of 
# applying the mask to the 'binary' vector
apply_mask <- function(mask, vec) {
  vec[which(mask == '1')] <- 1
  vec[which(mask == 'Z')] <- 0
  vec
}


input <- real_input  # The input
output <- list()     # List to hold output

# For each line in the input...
for (line in input) {
  type <- str_extract(line, '^\\w+')             # Get the 'type'
  if (type == 'mask') { masks <- to_mask(line)}  # 'mask' type to mask matrix
  if (type == 'mem') {
    # `address` is the number between the brackets
    address <- as.numeric(str_extract(line, '(?<=\\[)\\d+(?=\\])'))
    address_bits <- to_bits(address)
    value <- as.numeric(str_extract(line, '\\d+$'))  # Numbers after `=`
    
    # For each mask (row) in the mask matrix...
    for (r in seq(nrow(masks))) {
      # Mask the `address_bits`, convert it to a character representation of
      # the numeric value, then use that character representation as the name
      # of the list item to assign `value` to. Some of these list items
      # may be overwritten by subsequent iterations.
      masked_address_bits <- apply_mask(masks[r,], address_bits)
      masked_address <- as.character(to_numeric(masked_address_bits))
      output[[masked_address]] <- value
    }
  }
}
answer2 <- sum(unlist(output))