library(tidyverse)

# rotate matrix 90° clockwise
rotate <- function(mat, n = 1) {
    i <- 0
    while (i < n) {
        mat <- t(apply(mat, 2, rev))
        i <- i + 1
    }
    mat
}

# create a mirror image of the matrix
flip <- function(mat) {
    t(apply(mat, 1, rev))
}

# convert enhancement rule string into matrix
string_to_mat <- function(ii) {
    uu <- strsplit(ii, "/")[[1]]
    s <- nchar(uu[1])
    uu <- unlist(strsplit(uu, ""))
    matrix(uu, ncol = s, byrow = TRUE)
}

# convert matrix into enhancement rule string
mat_to_string <- function(m) {
    paste(apply(m, 1, paste, collapse = ""),
          collapse = "/")
}

# from an enhacement string, find all possible combinations
expand_combinations <- function(ii) {
    m <- string_to_mat(ii)
    s <- dim(m)[1]
    mr <- vector("list", 4)
    for (i in 0:3) {
        mr[[(i*2) + 1]] <- rotate(m, i + 1)        
        mr[[(i*2) + 2]] <- rotate(flip(m), i + 1)
    }
    map_chr(mr, mat_to_string)
}

# create data frame from input, and include all unique combinations
read_rules <- function(input) {
    read_delim(input, delim = "=",
               col_names = FALSE) %>%
        mutate_all(~ gsub(">?\\s+", "", .)) %>%
        set_names(c("ii", "oo")) %>%
        mutate(comb = map(ii, ~ expand_combinations(.))) %>%
        unnest() %>%
        distinct(comb, .keep_all = TRUE)
}

# split a matrix into a list that contains 2x2 or 3x3 
split_mat <- function(m, n) {
    n_mat <- dim(m)[1] / n
    res <- vector("list", n_mat^2)
    idx_start <- seq(1, to = dim(m)[1], by = n)
    k <- 1
    for (i in idx_start) {
        for (j in idx_start) {
            res[[k]] <- m[i:(i + n - 1), j:(j + n - 1)]
            k <- k + 1
        }
    }
    res
}

# reassamble a matrix from a list of 2x2 or 3x3 matrices
list_to_mat <- function(lst) {
    si <- nrow(lst[[1]])  
    s <- si * sqrt(length(lst))
    res <- array(, dim = c(s, s))
    idx_start <- seq(1, to = s, by = si)
    k <- 1
    for (i in idx_start) {
        for (j in idx_start) {
            res[i:(i + si - 1), j:(j + si - 1)] <- lst[[k]]
            k <- k + 1
        }
    }
    res
}

# use the enhancement rule book to grow the matrix
convert_rule <- function(pattern, rules) {
    res <- rules$oo[rules$comb == pattern]
    if (nchar(res) ==  0) stop("problem...")
    string_to_mat(res)       
}

# apply the enhacement algo for a set of rules, and n iterations
enhance <- function(rules, n) {
    i <- 0
    d <- c(2, 3)
    start <- c(".#./..#/###")
    m <- string_to_mat(start)

    while (i < n) {
        s <- dim(m)[1]
        sp <- d[s %% d == 0][1]
        smat <- split_mat(m, sp)
        sstr <- map(smat, mat_to_string)
        mlst <- map(sstr, convert_rule, rules)
        m <- list_to_mat(mlst)
        i <- i + 1
    }   
    m    
}

rules <- read_rules("advent-data/2017-12-21-data.txt")

part1 <- enhance(rules, 5)
sum(part1 == "#")

part2 <- enhance(rules, 18)
sum(part2 == "#")