Seating System

defmodule Aoc2020.Day11 do
  @moduledoc "Seating System"

  def run(), do: :timer.tc(fn -> part_2() end)

  def part_1() do
    parse_input()
    |> find_static_layout(&neighbour_seat_state/2)
    |> count_occupied_seats()
  end

  def part_2() do
    parse_input()
    |> find_static_layout(&visible_seat_state/2)
    |> count_occupied_seats()
  end

  def test_input do
    "L.LL.LL.LL\nLLLLLLL.LL\nL.L.L..L..\nLLLL.LL.LL\nL.LL.LL.LL\nL.LLLLL.LL\n..L.L.....\nLLLLLLLLLL\nL.LLLLLL.L\nL.LLLLL.LL"
  end

  def parse_input() do
    # test_input()
    File.read!("priv/inputs/2020/day11.txt")
    |> String.trim()
    |> String.split("\n")
    |> Enum.map(fn row ->
      String.codepoints(row)
      |> Enum.map(fn cell -> parse_value(cell) end)
      |> :array.from_list()
    end)
    |> :array.from_list()
  end

  def parse_value(str) do
    case str do
      "." -> :floor
      "L" -> :empty
      "#" -> :taken
    end
  end

  def get_cell(grid, x, y) do
    :array.get(x, :array.get(y, grid))
  end

  def get_neighbours(grid, x, y) do
    grid_height = :array.size(grid)
    grid_width = :array.size(:array.get(0, grid))

    neighbour_locations =
      for x2 <- (x - 1)..(x + 1),
          y2 <- (y - 1)..(y + 1),
          !(x2 == x && y2 == y),
          x2 >= 0 && x2 < grid_width,
          y2 >= 0 && y2 < grid_height,
          do: {x2, y2}

    Enum.map(neighbour_locations, fn {x, y} -> get_cell(grid, x, y) end)
  end

  def find_static_layout(grid, strategy) do
    next_state = evolve_seating(grid, strategy)

    if next_state == grid do
      grid
    else
      find_static_layout(next_state, strategy)
    end
  end

  def evolve_seating(grid, strategy) do
    :array.map(
      fn y, row ->
        :array.map(
          fn x, cell ->
            apply(strategy, [grid, {x, y, cell}])
          end,
          row
        )
      end,
      grid
    )
  end

  def neighbour_seat_state(grid, {_, _, :floor} = cell), do: :floor

  def neighbour_seat_state(grid, {x, y, :empty} = cell) do
    if count_occupied_neighbours(grid, cell) == 0, do: :taken, else: :empty
  end

  def neighbour_seat_state(grid, {x, y, :taken} = cell) do
    if count_occupied_neighbours(grid, cell) >= 4, do: :empty, else: :taken
  end

  def count_occupied_neighbours(grid, {x, y, _}) do
    get_neighbours(grid, x, y)
    |> Enum.filter(&(&1 == :taken))
    |> length()
  end

  def count_occupied_seats(grid) do
    :array.foldl(
      fn _, row, count ->
        count +
          :array.foldl(
            fn _, cell, count ->
              count + if cell == :taken, do: 1, else: 0
            end,
            0,
            row
          )
      end,
      0,
      grid
    )
  end

  # part 2
  def look_direction(grid, {x, y}, {xv, yv} = direction) do
    grid_height = :array.size(grid)
    grid_width = :array.size(:array.get(0, grid))
    x2 = x + xv
    y2 = y + yv

    if x2 >= 0 && x2 < grid_width && y2 >= 0 && y2 < grid_height do
      case get_cell(grid, x2, y2) do
        :taken -> :taken
        :empty -> :empty
        _ -> look_direction(grid, {x2, y2}, direction)
      end
    else
      :empty
    end
  end

  def count_visibly_occupied(grid, {x, y, _}) do
    directions = for xv <- [-1, 0, 1], yv <- [-1, 0, 1], {xv, yv} != {0, 0}, do: {xv, yv}

    Enum.map(directions, fn direction -> look_direction(grid, {x, y}, direction) end)
    |> Enum.filter(&(&1 == :taken))
    |> length()
  end

  def visible_seat_state(grid, {_, _, :floor} = cell), do: :floor

  def visible_seat_state(grid, {x, y, :empty} = cell) do
    if count_visibly_occupied(grid, cell) == 0, do: :taken, else: :empty
  end

  def visible_seat_state(grid, {x, y, :taken} = cell) do
    if count_visibly_occupied(grid, cell) >= 5, do: :empty, else: :taken
  end

  def render_to_string(grid) do
    IO.puts(
      :array.foldl(
        fn _, row, acc ->
          acc ++
            :array.foldl(
              fn _, cell, acc2 ->
                case cell do
                  :empty -> acc2 ++ 'L'
                  :taken -> acc2 ++ '#'
                  :floor -> acc2 ++ '.'
                  _ -> '?'
                end
              end,
              '',
              row
            ) ++ '\n'
        end,
        '',
        grid
      )
    )
  end
end