#!/usr/bin/env ruby

require "debug"
require "matrix"

@input = (ARGV.first.nil? ? DATA : ARGF)
  .readlines(chomp: true)
  .map(&:chars)
start = [@input.flatten.index { _1 == "S" }]
  .map { Vector[_1 / @input.size, _1 % @input.first.size] }
  .first
finish = [@input.flatten.index { _1 == "E" }]
  .map { Vector[_1 / @input.size, _1 % @input.first.size] }
  .first

NORTH = Vector[-1, 0]
EAST  = Vector[0, +1]
SOUTH = Vector[+1, 0]
WEST  = Vector[0, -1]
DIRECTIONS = [NORTH, EAST, SOUTH, WEST]

class PriorityQueue < Array
  alias extract_min shift
  alias peek first

  def <<(v)
    unless v.is_a?(Array)
      raise ArgumentError, "PriorityQueue requires array with priority as last"
    end

    super
    sort_by!(&:last)
  end
end

def dijkstra(graph, start, start_direction, target)
  dist = Hash.new # scores
  prev = Hash.new
  queue = PriorityQueue.new

  queue << [start, start_direction, 0]
  dist[[start, start_direction]] = 0
  # alternate only inserting source and add nodes inside score check later

  until queue.empty?
    position, direction, score = queue.extract_min

    neighbors = DIRECTIONS.each do |new_direction|
      neighbor = position + new_direction

      next if neighbor == position || graph.dig(*neighbor) == "#"

      new_score = score + 1 + (direction == new_direction ? 0 : 1) * 1_000

      if dist[[neighbor, new_direction]].nil? || new_score < dist[[neighbor, new_direction]]
        dist[[neighbor, new_direction]] = new_score
        prev[[neighbor, new_direction]] = [[position, direction]]

        queue << [neighbor, new_direction, new_score]
      elsif new_score == dist[[neighbor, new_direction]]
        prev[[neighbor, new_direction]] << [position, direction]
      end
    end
  end

  return dist, prev
end

def seats(dist, prev, target)
  visited, queue = Array.new, Array.new

  lowest_score = dist.keys.select { |(pos, dir)| pos == target }.map { dist[_1] }.min
  dist
    .select { |(position, _), score| position == target && lowest_score == score }
    .each { queue << _1.first }

  until queue.empty?
    directional_location = queue.pop
    visited << directional_location.first
    unless prev[directional_location].nil?
      queue.concat(prev[directional_location])
    end
  end

  return visited.uniq
end

dist, prev = dijkstra(@input, start, EAST, finish)
part_1 = dist.keys.select { |(pos, dir)| pos == finish }.map { dist[_1] }.min
p part_1
part_2 = seats(dist, prev, finish)
p part_2.size

__END__
###############
#.......#....E#
#.#.###.#.###.#
#.....#.#...#.#
#.###.#####.#.#
#.#.#.......#.#
#.#.#####.###.#
#...........#.#
###.#.#####.#.#
#...#.....#.#.#
#.#.#.###.#.#.#
#.....#...#.#.#
#.###.#.#.#.#.#
#S..#.....#...#
###############