defmodule Day10 do
  def part1(data) do
    grid = data
    |> String.split("\n", trim: true)
    |> Enum.map(&String.codepoints/1)
    |> Enum.map(fn v -> Enum.map(v, &String.to_integer/1) end)

    values = for {l, i} <- Enum.with_index(grid),
                 {v, j} <- Enum.with_index(l) do
               {v, {i, j}}
    end
    |> Enum.group_by(fn {v, _} -> v end, fn {_, p} -> p end)

    dist = fn {a, b}, {c, d} -> abs(a-c) + abs(b-d) end
    ends = fn p, prev ->
      for {_, v} <- Enum.filter(prev, fn {p2, _} -> dist.(p, p2) == 1 end),
          reduce: MapSet.new() do
        acc -> MapSet.union(acc, v)
      end
    end

    for i <- 8..0//-1,
        reduce: Map.get(values, 9, []) |> Enum.map(& {&1, MapSet.new([&1])}) do
      acc -> Map.get(values, i, [])
      |> Enum.map(fn p -> {p, ends.(p, acc)} end)
      |> Enum.filter(fn {_, m} -> MapSet.size(m) > 0 end)
    end
    |> Enum.map(fn {_, m} -> MapSet.size(m) end)
    |> Enum.sum()
  end

  def part2(data) do
    grid = data
    |> String.split("\n", trim: true)
    |> Enum.map(&String.codepoints/1)
    |> Enum.map(fn v -> Enum.map(v, &String.to_integer/1) end)

    values = for {l, i} <- Enum.with_index(grid),
                 {v, j} <- Enum.with_index(l) do
               {v, {i, j}}
    end
    |> Enum.group_by(fn {v, _} -> v end, fn {_, p} -> p end)

    acc = values
    |> Map.get(9, [])
    |> Enum.map(fn p -> {p, 1} end)

    dist = fn {a, b}, {c, d} -> abs(a-c) + abs(b-d) end
    count = fn p, prev ->
      prev
      |> Enum.filter(fn {p2, _} -> dist.(p, p2) == 1 end)
      |> Enum.map(fn {_, c} -> c end)
      |> Enum.sum()
    end

    for i <- 8..0//-1, reduce: acc do
      acc -> Map.get(values, i, [])
      |> Enum.map(fn p -> {p, count.(p, acc)} end)
      |> Enum.filter(fn {_, c} -> c > 0 end)
    end
    |> Enum.map(fn {_, c} -> c end)
    |> Enum.sum()
  end
end

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

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

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