# Setup ------------------------------------------------------------------------

test_input <- c(
  "sesenwnenenewseeswwswswwnenewsewsw",
  "neeenesenwnwwswnenewnwwsewnenwseswesw",
  "seswneswswsenwwnwse",
  "nwnwneseeswswnenewneswwnewseswneseene",
  "swweswneswnenwsewnwneneseenw",
  "eesenwseswswnenwswnwnwsewwnwsene",
  "sewnenenenesenwsewnenwwwse",
  "wenwwweseeeweswwwnwwe",
  "wsweesenenewnwwnwsenewsenwwsesesenwne",
  "neeswseenwwswnwswswnw",
  "nenwswwsewswnenenewsenwsenwnesesenew",
  "enewnwewneswsewnwswenweswnenwsenwsw",
  "sweneswneswneneenwnewenewwneswswnese",
  "swwesenesewenwneswnwwneseswwne",
  "enesenwswwswneneswsenwnewswseenwsese",
  "wnwnesenesenenwwnenwsewesewsesesew",
  "nenewswnwewswnenesenwnesewesw",
  "eneswnwswnwsenenwnwnwwseeswneewsenese",
  "neswnwewnwnwseenwseesewsenwsweewe",
  "wseweeenwnesenwwwswnew"
)
real_input <- readLines('input.txt')


# Mapping of compass directions to offsets in three-dimensional space
dir_map <- list(
  nw = c( 0,  1, -1),
  ne = c( 1,  0, -1),
  e =  c( 1, -1,  0),
  se = c( 0, -1,  1),
  sw = c(-1,  0,  1),
  w  = c(-1,  1,  0)
)

# Functions --------------------------------------------------------------------

# Given a vector of strings where each string represents a direction to move
# from the center `dirs` and a mapping of individual direction strings to 
# three-dimensional coordinate offsets `dir_map`, return the final offset
# in three-dimensional coordinates indicated by the directions.
parse_directions <- function(dirs, dir_map) {
  final_location <- c(0, 0, 0)
  for (dir in dirs) { final_location <- final_location + dir_map[[dir]] }
  final_location
}


# Given a vector of strings from the puzzle input `input` and a mapping of
# individual compass directions to three-dimensional coordinate offsets 
# `dir_map`, return a list of three-dimensional offsets resulting from following
# the set of directions in each element of `input`
parse_input <- function(input, dir_map) {
  tokens <- regmatches(input, gregexpr('(e|w|[ns][ew])', input))
  lapply(tokens, parse_directions, dir_map)
}


# Given a the maximum distance from a central point of any individual tile 
# offset `max_offset`, return a three-dimensional array large enough to 
# encompass the entire 'floor', where each array element is either `NA` or `0`. 
# A `0` represents the location of a white tile in the floor, an `NA` represents 
# a space in the matrix that is not a part of the floor. These `NA`'s are 
# essentially ignored.
init_tile_map <- function(max_offset) {
  offset_range <- c(-max_offset:max_offset)     # The maximum range of offsets
  map_dim <- (max_offset * 2) + 1               # The dimensions of the array
  tile_map <- array(NA, dim = rep(map_dim, 3))  # Empty array of sufficient size
  center <- rep(ceiling(map_dim / 2), 3)        # The center point of the array
  
  # Only the array indices that sum to the same value as the central index 
  # would be included in the diagonal slice of the array representing the
  # floor. Those tiles start off as white (`0`)
  all_coords <- arrayInd(1:map_dim^3, rep(map_dim, 3))
  tile_coords <- all_coords[rowSums(all_coords) == sum(center),]
  tile_map[tile_coords] <- 0
  
  tile_map  # Return the tile_map
}


# Given an offset from the `tile_map`'s central point `relative_tile_loc` and 
# a three-dimensional array containing the locations of tiles in the floor 
# `tile_map`, 'flips' the tile indicated by `relative_tile_loc`, i.e. toggles
# the array index value between `1` and `0`
flip_tile <- function(relative_tile_loc, tile_map) {
  center <- ceiling(dim(tile_map) / 2)  # Array central point
  ti <- center + relative_tile_loc      # Absolute array index of the tile
  
  # 'Flip' the tile (toggle the value)
  tile_map[ti[1], ti[2], ti[3]] <- abs(1 - tile_map[ti[1], ti[2], ti[3]])

  tile_map  # Return the modified matrix
}


# Processing -------------------------------------------------------------------

relative_tile_locs <- parse_input(real_input, dir_map)  # Parse the input directions
max_offset <- max(abs(unlist(relative_tile_locs)))      # Max tile offset in any direction
tile_map <- init_tile_map(max_offset)                   # Initialize the tile map array

# For each tile that needs to be flipped, flip it!
for (loc in relative_tile_locs) { tile_map <- flip_tile(loc, tile_map) }

answer1 <- sum(tile_map, na.rm = T)  # Sum of all tile values, black tiles = 1

# Setup ------------------------------------------------------------------------

test_input <- c(
  "sesenwnenenewseeswwswswwnenewsewsw",
  "neeenesenwnwwswnenewnwwsewnenwseswesw",
  "seswneswswsenwwnwse",
  "nwnwneseeswswnenewneswwnewseswneseene",
  "swweswneswnenwsewnwneneseenw",
  "eesenwseswswnenwswnwnwsewwnwsene",
  "sewnenenenesenwsewnenwwwse",
  "wenwwweseeeweswwwnwwe",
  "wsweesenenewnwwnwsenewsenwwsesesenwne",
  "neeswseenwwswnwswswnw",
  "nenwswwsewswnenenewsenwsenwnesesenew",
  "enewnwewneswsewnwswenweswnenwsenwsw",
  "sweneswneswneneenwnewenewwneswswnese",
  "swwesenesewenwneswnwwneseswwne",
  "enesenwswwswneneswsenwnewswseenwsese",
  "wnwnesenesenenwwnenwsewesewsesesew",
  "nenewswnwewswnenesenwnesewesw",
  "eneswnwswnwsenenwnwnwwseeswneewsenese",
  "neswnwewnwnwseenwseesewsenwsweewe",
  "wseweeenwnesenwwwswnew"
)
real_input <- readLines('input.txt')


# Mapping of compass directions to offsets in three-dimensional space
dir_map <- list(
  nw = c( 0,  1, -1),
  ne = c( 1,  0, -1),
  e =  c( 1, -1,  0),
  se = c( 0, -1,  1),
  sw = c(-1,  0,  1),
  w  = c(-1,  1,  0)
)


# Functions --------------------------------------------------------------------

# Given a vector of strings where each string represents a direction to move
# from the center `dirs` and a mapping of individual direction strings to 
# three-dimensional coordinate offsets `dir_map`, return the final offset
# in three-dimensional coordinates indicated by the directions.
parse_directions <- function(dirs, dir_map) {
  final_location <- c(0, 0, 0)
  for (dir in dirs) { final_location <- final_location + dir_map[[dir]] }
  final_location
}


# Given a vector of strings from the puzzle input `input` and a mapping of
# individual compass directions to three-dimensional coordinate offsets 
# `dir_map`, return a list of three-dimensional offsets resulting from following
# the set of directions in each element of `input`
parse_input <- function(input, dir_map) {
  tokens <- regmatches(input, gregexpr('(e|w|[ns][ew])', input))
  lapply(tokens, parse_directions, dir_map)
}


# Given a the maximum distance from a central point of any individual tile 
# offset `max_offset`, returns a data frame containing columns for `x`, `y`, 
# `z`, and `color`, where each row in the data frame represents an individual
# tile with coordinates.
build_tile_table <- function(max_offset) {
  offset_range <- c(-(max_offset):(max_offset))   # The maximum range of offsets
  
  # Create a data frame containing all combinations of the values in the
  # `offset_range`, then filter out any rows where the coordinates do not
  # sum to zero
  all_coords <- expand.grid(x = offset_range, y = offset_range, z = offset_range)
  tile_coords <- all_coords[rowSums(all_coords) == 0,]
  
  # Return a data frame containing all the remaining coordinates and 'white'
  # for the color
  data.frame(
    x = tile_coords$x, y = tile_coords$y, z = tile_coords$z, 
    color = 'white', stringsAsFactors = FALSE
  )
}


# Given an index relative to the central tile `relative_tile_loc` and a data
# frame of tile locations, 'flip' the color of the tile on the row indicated
# by `relative_tile_loc`
flip_tile <- function(relative_tile_loc, tile_table) {
  selector <- (
    tile_table$x == relative_tile_loc[1] &
    tile_table$y == relative_tile_loc[2] &
    tile_table$z == relative_tile_loc[3]
  )
  current_color <- tile_table[selector, 'color']
  tile_table[selector, 'color'] <- ifelse(current_color == 'white', 'black', 'white')
  tile_table
}


# Processing -------------------------------------------------------------------

relative_tile_locs <- parse_input(real_input, dir_map)  # Parse the input directions
max_offset <- max(abs(unlist(relative_tile_locs)))      # Max tile offset in any direction
tile_table <- build_tile_table(max_offset)              # Create a data frame of tiles

# For each set of directions, flip the indicated tile
for (loc in relative_tile_locs) { tile_table <- flip_tile(loc, tile_table) }
answer1 <- sum(tile_table$color == 'black')  # Count the black tiles

# Setup ------------------------------------------------------------------------

source('exercise_1.R')


# Function ---------------------------------------------------------------------

# Given a index to a three-dimensional array `i` and a three-dimensional array
# `tile_map`, returns the value at that index of the array. Provides safety 
# against referencing indices that aren't actually present in the array.
get_tile_value <- function(i, tile_map) {
  # If the index is not in the array, `val` is NA
  val <- tryCatch(
    val <- tile_map[i[1], i[2], i[3]],
    error = function(e) { NA_real_ }
  )
  
  # Oddly, if any of the elements of `i` is `0`, the index operator doesn't
  # throw an error but returns numeric(0), so we need to test for that.
  if (length(val) > 0) { val } else { NA_real_ }
}


# Given an index to a three-dimensional array `tile_index` and a 
# three-dimensional array `tile_map`, returns the number of black tiles 
# neighboring the element at index `tile_index`.
get_neighbors <- function(tile_index, tile_map) {
  arr_index <- arrayInd(tile_index, dim(tile_map))

  neighbor_offsets <- list(
    c(0,  1, -1), c( 1, 0, -1), c( 1, -1, 0),
    c(0, -1,  1), c(-1, 0,  1), c(-1,  1, 0)
  )
  neighbor_indices <- lapply(neighbor_offsets, `+`, arr_index)
  neighbor_values <- vapply(neighbor_indices, get_tile_value, numeric(1), tile_map)
  sum(neighbor_values, na.rm = T)
}


# Given a 3D array `tile_map`, returns the elements (in a 1D vector) on the
# outer edges of the array.
get_shell <- function(tile_map) {
  d <- dim(tile_map)
  d1_edge_cells <- tile_map[c(1, d[1]), ,  ]
  d2_edge_cells <- tile_map[ , c(1, d[2]), ]
  d3_edge_cells <- tile_map[ , , c(1, d[3])]
  c(d1_edge_cells, d2_edge_cells, d3_edge_cells)
}


# Given an index to a 3D array `tile_index` and a 3D array `tile_map`, return
# the state of the element at index `tile_index` after applying the rules
# for changing element state in the puzzle description.
next_tile_state <- function(tile_index, tile_map) {
  tile_value <- tile_map[tile_index]
  if (is.na(tile_value)) { return(NA_real_) }
  
  neighbors_value <- get_neighbors(tile_index, tile_map)
  if (tile_value == 1 && (neighbors_value == 0 | neighbors_value > 2)) {
    0
  } else if (tile_value == 0 && neighbors_value == 2) {
    1
  } else {
    tile_value
  }
}


# Given a 3D array `tile_map` and the number of 'layers' to expand the array by
# `expand_by`, add `expand_by` additional layers to the outside of the array
# and return it.
expand_tile_map <- function(tile_map, expand_by = 1) {
  old_dims <- dim(tile_map)
  new_dims <- old_dims + (2 * expand_by)
  new_center <- ceiling(new_dims / 2)
  new_map <- array(dim = new_dims)
  
  all_coords <- arrayInd(1:prod(new_dims), new_dims)
  tile_coords <- all_coords[rowSums(all_coords) == sum(new_center),]
  new_map[tile_coords] <- 0
  
  r1 <- new_center - floor(old_dims / 2)
  r2 <- new_center + floor(old_dims / 2)
  new_map[r1[1]:r2[1], r1[2]:r2[2], r1[3]:r2[3]] <- tile_map
  
  new_map
}

# Processing -------------------------------------------------------------------

# Advance the floor state 100 times, starting with the floor state at the end
# of part one.
rounds <- 100
pb <- txtProgressBar(max = rounds, style = 3)
for (i in 1:rounds) {
  
  # If there are any black tiles on the outer edge of the `tile_map`, then we 
  # 'may' need to flip tiles that don't exist yet, so we go ahead and add an
  # extra layer of white tiles to the outside.
  if (any(get_shell(tile_map) == 1, na.rm = T)) { 
    tile_map <- expand_tile_map(tile_map) 
  }
  
  # Advance the floor state by iterating over the `tile_map` and calculating 
  # the next state for each element
  tile_map <- apply(
    slice.index(tile_map, c(1, 2, 3)), 
    c(1, 2, 3), 
    next_tile_state, 
    tile_map
  )
  
  setTxtProgressBar(pb, i)
}
close(pb)

answer2 <- sum(tile_map, na.rm = T)  # Sum the number of black tiles

# Setup ------------------------------------------------------------------------

source('exercise_1b.R')


# Functions --------------------------------------------------------------------

# Given a vector indicating the x/y/z coordinates of a tile `tile_loc` and a
# data frame of tile locations and colors `tile_table`, return the color of
# the tile on the row indicated by `tile_loc`
get_tile_color <- function(tile_loc, tile_table) {
  selector <- (
    tile_table$x == tile_loc[1] &
      tile_table$y == tile_loc[2] &
      tile_table$z == tile_loc[3]
  )
  tile_table[selector, 'color']
}


# Given a vector indicating the x/y/z coordinates of a tile `tile_loc` and a
# data frame of tile locations and colors `tile_table`, returns a list of 
# coordinates for neighboring tiles
get_neighbors <- function(tile_loc, tile_table) {
  neighbor_offsets <- list(
    c(0,  1, -1), c( 1, 0, -1), c( 1, -1, 0),
    c(0, -1,  1), c(-1, 0,  1), c(-1,  1, 0)
  )
  lapply(neighbor_offsets, `+`, tile_loc)
}


# Given a vector indicating the x/y/z coordinates of a tile `tile_loc` and a
# data frame of tile locations and colors `tile_table`, return a vector of the
# colors of the neighboring tiles
get_neighbor_colors <- function(tile_loc, tile_table) {
  neighbor_indices <- get_neighbors(tile_loc, tile_table)
  neighbor_colors <- sapply(neighbor_indices, get_tile_color, tile_table)
  neighbor_colors
}


# Given a vector indicating the x/y/z coordinates of a tile `tile_loc` and a
# data frame of tile locations and colors `tile_table`, return the next color
# of the indicated tile according to the rules of the puzzle
next_tile_color <- function(tile_loc, tile_table) {
  tile_color <- get_tile_color(tile_loc, tile_table)  # Current tile color
  
  # Colors of the neighboring tiles
  neighbor_colors <- get_neighbor_colors(tile_loc, tile_table)
  black_neighbors <- sum(neighbor_colors == 'black', na.rm = T)  # Count black neighbors

  # Should the tile be flipped?
  flip <- if (tile_color == 'black') {
    black_neighbors == 0 | black_neighbors > 2
  } else if (tile_color == 'white') {
    black_neighbors == 2
  } else {
    stop(paste('Could not get a color for tile: ', tile_loc))
  }
  
  # The final tile color
  if (flip & tile_color == 'white') { 
    'black'
  } else if (flip & tile_color == 'black') {
    'white'
  } else {
    tile_color
  }
}


# Given a vector indicating the x/y/z coordinates of a tile `tile_loc` and a
# data frame of tile locations and colors `tile_table`, add rows to the 
# `tile_table` to ensure that the `tile_table` contains records for all the
# tiles neighboring `tile_loc`
add_neighbors <- function(tile_loc, tile_table) {
  neighbor_indices <- get_neighbors(tile_loc, tile_table)  # Locations of neighbors
  
  # For each neighbor index...
  for (ni in neighbor_indices) {
    
    # Define a selector for the row at the index
    ni <- as.numeric(ni)
    selector <- (
      tile_table$x == ni[1] &
        tile_table$y == ni[2] &
        tile_table$z == ni[3]
    )
    
    # If there's no row at that index, add one
    if (nrow(tile_table[selector,]) < 1) {
      new_row <- data.frame(
        x = ni[1], y = ni[2], z = ni[3], 
        color = 'white', all_neighbors = FALSE,
        stringsAsFactors = F
      )
      tile_table <- rbind(tile_table, new_row)
    }
  }
  
  # Indicate in the 'all_neighbors' column that all the neighbors for this
  # tile are represented in the table
  tile_table[
    tile_table$x == tile_loc[1] &
    tile_table$y == tile_loc[2] &
    tile_table$z == tile_loc[3],
    'all_neighbors'
  ] <- TRUE
  
  tile_table  # Return the `tile_table`
}

# Given a data frame of tile locations and colors `tile_table`, iterate through
# the black tiles and ensure those tiles all have their neighbors represented
# in the table
expand_table <- function(tile_table) {
  
  # Select the coordinates from all the `tile_table` rows representing black 
  # tiles where we haven't yet confirmed all the neighbors for that  tile are
  # in the `tile_table`
  black_tile_locs <- tile_table[
    (tile_table$color == 'black' & !tile_table$all_neighbors), 
    c('x', 'y', 'z')
  ]
  num_black_tiles <- nrow(black_tile_locs)  # The number of tiles to check
  
  # For each black tile to be checked...
  for (btl in split(black_tile_locs, 1:num_black_tiles)) {
    tile_table <- add_neighbors(btl, tile_table)  # Add all its neighbors
  }
  tile_table  # Return the `tile_table`
}


# Given a data frame of tile locations and colors `tile_table`, return the 
# colors for each row in the data frame representing the next state of each
# tile according to the instructions
next_table_colors <- function(tile_table) {
  apply(tile_table[,c('x', 'y', 'z')], 1, next_tile_color, tile_table)
}


# Processing -------------------------------------------------------------------

tile_table$all_neighbors <- FALSE  # Set the `all_neighbors` column to FALSE

# Advance the floor state 100 times, starting with the floor state at the end
# of part one.
rounds <- 100
pb <- txtProgressBar(min = 0, max = rounds, style = 3)
for (i in 1:rounds) {
  tile_table <- expand_table(tile_table)       # Ensure all black tiles have neighbors
  new_colors <- next_table_colors(tile_table)  # Determine new tile colors
  tile_table[, 'color'] <- new_colors          # Set new tile colors
  setTxtProgressBar(pb, i)
}
close(pb)

answer2 <- sum(tile_table[,'color'] == 'black')  # Count the black tiles