defmodule Day15 do
  defmodule Game do
    @derive Inspect
    defstruct [:grid, :robot, :moves]

    def from_data(data) do
      [grid, moves] = data
      |> String.split("\n\n", trim: true)

      grid = grid
      |> String.split("\n", trim: true)
      |> Enum.with_index()
      |> Enum.flat_map(fn {s, i} ->
        String.codepoints(s)
        |> Enum.with_index()
        |> Enum.map(fn {c, j} -> {{i, j}, c} end)
      end)
      |> Map.new()

      robot = grid
      |> Enum.find_value(fn {pos, c} -> c == "@" and pos end)

      moves = moves
      |> String.split("\n", trim: true)
      |> Enum.flat_map(&String.codepoints/1)

      %Game{grid: grid, robot: robot, moves: moves}
    end

    def expand(game) do
      grid = game.grid
      |> Enum.flat_map(fn {{i, j}, v} ->
        case v do
          "." -> [{{i, j*2}, "."}, {{i, j*2+1}, "."}]
          "@" -> [{{i, j*2}, "@"}, {{i, j*2+1}, "."}]
          "O" -> [{{i, j*2}, "["}, {{i, j*2+1}, "]"}]
          "#" -> [{{i, j*2}, "#"}, {{i, j*2+1}, "#"}]
        end
      end)
      |> Map.new()

      robot = {elem(game.robot, 0), elem(game.robot, 1)*2}

      %Game{grid: grid, robot: robot, moves: game.moves}
    end

    def do_moves(game) do
      if game.moves == [] do
        game
      else
        move(game)
        |> do_moves()
      end
    end

    defp move(game) do
      [move | moves] = game.moves

      moved? = check_moves(game.grid, move, game.robot) |> elem(0)

      grid =
        if moved? do
          apply_move(game.grid, move, queue_new(game.robot))
        else
          game.grid
        end

      robot = if moved? do next_pos(game.robot, move) else game.robot end

      %Game{grid: grid, robot: robot, moves: moves}
    end

    defp check_moves(grid, move, k, vis \\ MapSet.new()) do
      if MapSet.member?(vis, k) do
        {true, vis}
      else
        vis = MapSet.put(vis, k)
        v = Map.fetch!(grid, k)
        case v do
          "." ->
            {true, vis}
          "#" ->
            {false, vis}
          "@" ->
            check_moves(grid, move, next_pos(k, move), vis)
          "O" ->
            check_moves(grid, move, next_pos(k, move), vis)
          "[" ->
            nk = next_pos(k, move)
            if move == "^" or move == "v" do
              {oi, _} = k
              {_, nj} = nk
              {a1, vis} = check_moves(grid, move, nk, vis)
              {a2, vis} = check_moves(grid, move, {oi, nj+1}, vis)
              {a1 and a2, vis}
            else
              check_moves(grid, move, nk, vis)
            end
          "]" ->
            nk = next_pos(k, move)
            if move == "^" or move == "v" do
              {oi, _} = k
              {_, nj} = nk
              {a1, vis} = check_moves(grid, move, nk, vis)
              {a2, vis} = check_moves(grid, move, {oi, nj-1}, vis)
              {a1 and a2, vis}
            else
              check_moves(grid, move, nk, vis)
            end
        end
      end
    end

    defp apply_move(grid, move, queue) do
      if queue_empty(queue) do
        grid
      else
        {pos, nv, queue} = queue_pop(queue, move)
        v = Map.fetch!(grid, pos)
        case v do
          "#" ->
            raise "Unexpected"
          "." ->
            apply_move(Map.put(grid, pos, nv), move, queue)
          v ->
            grid = Map.put(grid, pos, nv)
            npos = next_pos(pos, move)
            nv = Map.fetch!(grid, npos)
            queue =
              if (move == "^" or move == "v") and (nv == "[" or nv == "]") do
                {i, j} = npos
                j = if nv == "[" do j+1 else j-1 end
                queue = queue_add(queue, npos, v)
                queue_add(queue, {i, j}, ".")
              else
                queue_add(queue, npos, v)
              end
            apply_move(grid, move, queue)
        end
      end
    end

    defp queue_add(queue, pos, v) do
      case :gb_trees.lookup(pos, queue) do
        :none ->
          :gb_trees.insert(pos, v, queue)
        {_, "."} ->
          :gb_trees.update(pos, v, queue)
        _ ->
          queue
      end
    end

    defp queue_pop(queue, move) do
      if move == "^" do
        :gb_trees.take_largest(queue)
      else
        :gb_trees.take_smallest(queue)
      end
    end

    defp queue_new(robot) do
      :gb_trees.insert(robot, ".", :gb_trees.empty())
    end

    defp queue_empty(queue) do
      :gb_trees.is_empty(queue)
    end

    defp next_pos({i, j}, move) do
      case move do
        ">" -> {i, j+1}
        "^" -> {i-1, j}
        "<" -> {i, j-1}
        "v" -> {i+1, j}
      end
    end

    def print(game) do
      w = game.grid |> Enum.map(fn {{_, j}, _} -> j end) |> Enum.max()
      h = game.grid |> Enum.map(fn {{i, _}, _} -> i end) |> Enum.max()
      for i <- 0..h do
        for j <- 0..w, reduce: [] do
          acc -> [Map.fetch!(game.grid, {i, j}) | acc]
        end
        |> Enum.reverse()
        |> Enum.join("")
        |> IO.puts()
      end
      game
    end
  end

  def part1(data) do
    data
    |> Game.from_data()
    |> Game.do_moves()
    |> then(& &1.grid)
    |> Enum.filter(fn {_, v} -> v == "O" end)
    |> Enum.map(fn {{i, j}, _} -> i*100 + j end)
    |> Enum.sum()
  end

  def part2(data) do
    data
    |> Game.from_data()
    |> Game.expand()
    |> Game.do_moves()
    |> then(& &1.grid)
    |> Enum.filter(fn {_, v} -> v == "[" end)
    |> Enum.map(fn {{i, j}, _} -> i*100 + j end)
    |> Enum.sum()
  end
end

data = IO.read(:stdio, :eof)

{time1, ans1} = :timer.tc(fn -> Day15.part1(data) end)
IO.puts("Time  : #{time1 / 1000000}")
IO.puts("Answer: #{ans1}")

{time2, ans2} = :timer.tc(fn -> Day15.part2(data) end)
IO.puts("Time  : #{time2 / 1000000}")
IO.puts("Answer: #{ans2}")