#!/usr/bin/env ruby

require "debug"
require "matrix"

input = (ARGV.first.nil? ? DATA : ARGF)
  .readlines(chomp: true)
input2 = input
  .map { _1.gsub("#", "##").gsub("O", "[]").gsub(".", "..").gsub("@", "@.") }

split = input.index("")

DIRECTIONS = {
  "^" => Vector[-1, 0],
  ">" => Vector[0, 1],
  "v" => Vector[1, 0],
  "<" => Vector[0, -1]
}

@grid = input[0...split].map(&:chars)
@grid2 = input2[0...split].map(&:chars)
@moves = input[split+1..].map(&:chars).flatten

def find_robot
  @robot = nil
  (0...@grid.size).each do |x|
    break unless @robot.nil?
    (0...@grid.first.size).each do |y|
      if @grid[x][y] == "@"
        @grid[x][y] = "."
        return Vector[x, y]
      end
    end
  end
end

def print_grid(grid)
  counts = Hash.new { |h, k| h[k] = 0 }
  grid.each_with_index do |row, x|
    row.each_with_index do |val, y|
      if @robot == Vector[x, y]
        counts["@"] += 1
        print "@"
      else
        counts[val] += 1
        print val
      end
    end
    puts
  end
  p counts
  p @robot
  puts
  puts
end

def push(location, direction, called_by, test:)
  return false if @grid.dig(*location) == "#"
  return true if @grid.dig(*location) == "."

  location_value = @grid.dig(*location)

  move_to = location + direction

  result = if ["[", "]"].include?(location_value)
    move_to = location + direction
    pair = location + (location_value == "[" ? Vector[0, 1] : Vector[0, -1])

    if pair == move_to || pair == called_by
      push(move_to, direction, location, test: test)
    else
      push(move_to, direction, location, test: test) && push(pair, direction, location, test: test)
    end
  else
    push(move_to, direction, location, test: test)
  end

  if result && test == false
    @grid[move_to[0]][move_to[1]] = location_value
    @grid[location[0]][location[1]] = "."
  end

  return result
end

def predict
  @moves.each do |move|
    robot_next = @robot + DIRECTIONS[move]

    block = @grid.dig(*robot_next.to_a)

    case block
    when "#" # robot does not move
    when "." # robot can take position
      @robot = robot_next
    else # it's a box (of some sort)
      res = push(robot_next, DIRECTIONS[move], @robot, test: true)
      if res
        push(robot_next, DIRECTIONS[move], @robot, test: false)
        @robot = robot_next
      end
    end
  end
end

def compute_gps
  (0...@grid.size).map do |x|
    (0...@grid.first.size).map do |y|
      ["O", "["].include?(@grid[x][y]) ? 100 * x + y : 0
    end
  end.flatten.sum
end

# part 1
@robot = find_robot
predict
p compute_gps

# part 2
@grid = @grid2
@robot = find_robot
predict
p compute_gps

__END__
######
#....#
#@O..#
#.OO.#
#.O..#
#.O..#
#.#..#
#....#
######

>>^>>v