aoc2022/day9/day9.jl

#!/usr/bin/env julia

using Test
using Images

struct Move
  direction :: Char
  count :: Int
end

mutable struct Point{T}
  x::T
  y::T
end

function parse_move(line)
  parts = split(line, limit=2)
  return Move(parts[1][1], parse(Int, parts[2]))
end

function read_moves(io)
  return [parse_move(line) for line in readlines(io)]
end

function apply_move(head, tail, direction)
  if direction == 'U'
    head.y += 1
  elseif direction == 'D'
    head.y -= 1
  elseif direction == 'L'
    head.x -= 1
  elseif direction == 'R'
    head.x += 1
  elseif direction != 'O'
    throw(DomainError("Unknown direction"))
  end
  x_diff = head.x - tail.x
  y_diff = head.y - tail.y
  if abs(x_diff) <= 1 && abs(y_diff) <= 1
    return head, tail
  end
  if abs(x_diff) <= 1
    tail.x += x_diff
    tail.y += sign(y_diff) * (abs(y_diff) - 1)
  elseif abs(y_diff) <= 1
    tail.y += y_diff
    tail.x += sign(x_diff) * (abs(x_diff) - 1)
  elseif abs(x_diff) == 2 && abs(y_diff) == 2
    # possible with more than 2 knots
    tail.x += sign(x_diff) * (abs(x_diff) - 1)
    tail.y += sign(y_diff) * (abs(y_diff) - 1)
  else
    throw(DomainError("invalid knots " * string(head) * " " * string(tail)))
  end 
  return head, tail
end

function get_visited(moves; nknots::Int=2)
  visited = Set{NTuple{2, Int}}()
  knots = [Point(0, 0) for _ in 1:nknots]
  push!(visited, (knots[nknots].x, knots[nknots].y))
  for move in moves
    for _ in 1:move.count
      knots[1], knots[2] = apply_move(knots[1], knots[2], move.direction)
      for i in 2:nknots
        _, knots[i] = apply_move(knots[i-1], knots[i], 'O')
      end
      push!(visited, (knots[nknots].x, knots[nknots].y))
    end
  end
  return visited
end

function int_map(visited)
  (min_x, max_x) = extrema([t[1] for t in visited])
  (min_y, max_y) = extrema([t[2] for t in visited])
  nrows = max_y - min_y + 1
  ncols = max_x - min_x + 1
  imap = zeros(Int, nrows, ncols)
  for (x, y) in visited
    imap[nrows - (y - min_y), (x - min_x + 1)] = 1
  end
  imap[nrows + min_y, 1 - min_x] = -1
  return imap
end

function ascii_map(visited)
  map_char(n) = (n == -1 ? 's'
                 : (n == 1 ? "#"[1] : '.'))
  srows = mapreduce(map_char, *, int_map(visited), dims=2; init="")
  return join(srows, "\n")
end

function scale_pixels(img, scale::Int)
  (nrows, ncols) = size(img)
  new_img = Matrix{RGB}(undef, nrows * scale, ncols * scale)
  for i in 1:nrows
    for j in 1:ncols
      i2 = (i - 1) * scale + 1
      j2 = (j - 1) * scale + 1
      new_img[i2:i2+scale-1, j2:j2+scale-1] .= img[i, j]
    end
  end
  return new_img
end

function image_map(visited; scale::Int=10)
  imap = int_map(visited)
  map_color(n) = (n == -1 ? RGB(0, 0.8, 0)
                  : (n == 1 ? RGB(0.2, 0.2, 0.0) : RGB(0, 0, 0)))
  cmap = map(map_color, int_map(visited))
  return scale_pixels(cmap, scale)
end

function test()
  @testset "visited count 2 knot" verbose=true begin
    @test length(get_visited(read_moves("example.txt"))) == 13
    @test length(get_visited(read_moves("input.txt"))) == 6037
  end
  @testset "visited count 10 knot" verbose=true begin
    @test length(get_visited(read_moves("example.txt"), nknots=10)) == 1
    @test length(get_visited(read_moves("input.txt"), nknots=10)) == 2485
  end
end

function main()
  if size(ARGS, 1) == 0
    test()
  else
    infile = ARGS[1]
    println("infile = ", infile)
    moves = read_moves(infile)
    println("moves: ", moves)
    visited = get_visited(moves)
    println("visited: ", visited)
    println("count  : ", length(visited))
    println(ascii_map(visited))
    save("visited.png", image_map(visited))

    visited10 = get_visited(moves, nknots=10)
    println("visited10: ", visited10)
    println("count10  : ", length(visited10))
    println(ascii_map(visited10))
    save("visited10.png", image_map(visited10))
  end
end

main()