Day 14

Mix.install([
  {:kino_vega_lite, "~> 0.1.7"}
])

alias VegaLite, as: Vl

Section

defmodule Day14 do
  def part1(input) do
    bottom = input |> Map.keys() |> Enum.max_by(&elem(&1, 1)) |> elem(1)

    input
    |> Stream.unfold(&fall(&1, bottom, {500, 0}))
    |> Enum.reduce({nil, 0}, fn state, {_, count} ->
      {state, count + 1}
    end)
  end

  def part2(input) do
    bottom = input |> Map.keys() |> Enum.max_by(&elem(&1, 1)) |> elem(1) |> Kernel.+(2)

    state =
      (-bottom - 1)..(bottom + 1)
      |> Stream.map(&{500 + &1, bottom})
      |> Stream.map(&{&1, :Rock})
      |> Enum.into(input)

    state
    |> Stream.unfold(&fall(&1, bottom, {500, 0}))
    |> Enum.reduce({nil, 0}, fn state, {_, count} ->
      {state, count + 1}
    end)
  end

  defp fall(state, bottom, {x, y} = sand) do
    if state[{500, 0}] do
      nil
    else
      air = Enum.find([{x, y + 1}, {x - 1, y + 1}, {x + 1, y + 1}], &(!state[&1]))

      case air do
        nil -> {state, Map.put(state, sand, :Sand)}
        {_, y} when y > bottom -> nil
        _ -> fall(state, bottom, air)
      end
    end
  end

  def parse_input(path) do
    path
    |> File.stream!()
    |> Stream.flat_map(&parse_line/1)
    |> Map.new(fn coord -> {coord, :Rock} end)
  end

  defp parse_line(line) do
    line
    |> String.split(~r/\D+/, trim: true)
    |> Stream.map(&String.to_integer/1)
    |> Stream.chunk_every(2)
    |> Stream.chunk_every(2, 1, :discard)
    |> Stream.flat_map(fn [[a, b], [c, d]] ->
      for x <- a..c, y <- b..d, do: {x, y}
    end)
  end

  def visualize(state, opts \\ []) do
    Vl.new(opts)
    |> Vl.data_from_values(for {{x, y}, type} <- state, do: %{"x" => x, "y" => y, "type" => type})
    |> Vl.mark(:rect, tooltip: true)
    |> Vl.encode_field(:x, "x", axis: nil)
    |> Vl.encode_field(:y, "y", axis: nil)
    |> Vl.encode(:color, field: "type", type: :nominal)
  end
end

input = Day14.parse_input("#{__DIR__}/day14.txt")

{state, count} = Day14.part1(input)

count

Day14.visualize(state, width: 500, height: 1000)

{state, count} = Day14.part2(input)

count

Day14.visualize(state, width: 600, height: 400)