aoc2017_day7.ex

defmodule Day7 do
  def part1(), do:
    root(read_programs_map()).name

  def part2() do
    programs_map = read_programs_map()
    {:error, {:unbalanced, unbalanced_info}} = total_weight(programs_map, root(programs_map))
    unbalanced_info.program.weight + unbalanced_info.missing_weight
  end

  defp root(programs_map) do
    possible_roots = programs_map |> Map.values() |> Stream.map(&(&1.name)) |> MapSet.new()
    all_children = programs_map |> Map.values() |> Stream.flat_map(&(&1.children)) |> MapSet.new()
    [root_name] = MapSet.difference(possible_roots, all_children) |> Enum.to_list()
    program(programs_map, root_name)
  end

  defp program(programs_map, program_name), do:
    Map.fetch!(programs_map, program_name)

  defp children(programs_map, program), do:
    Stream.map(program.children, &program(programs_map, &1))

  defp total_weight(programs_map, program) do
    with \
      children = children(programs_map, program),
      total_weight_results = Enum.map(children, &total_weight(programs_map, &1)),
      {:ok, children_total_weights} <- extract_and_verify_results(total_weight_results),
      :ok <- verify_balance(children, children_total_weights),
      do: {:ok, program.weight + Enum.sum(children_total_weights)}
  end

  defp extract_and_verify_results(total_weight_results) do
    case Enum.split_with(total_weight_results, &match?({:error, _}, &1)) do
      {[], ok_weights} -> {:ok, Enum.map(ok_weights, fn {:ok, weight} -> weight end)}
      {[first_error | _other_errors], _ok_weights} -> first_error
    end
  end

  defp verify_balance(programs, total_weights) do
    case \
      total_weights
      |> Stream.zip(programs)
      |> Enum.group_by(fn {weight, _program} -> weight end, fn {_weight, program} -> program end)
      |> Enum.sort_by(fn {_weight, programs} -> length(programs) end)
    do
      [{unbalanced_weight, [unbalanced_program]}, {correct_weight, [_, _ | _] = _correct_programs}] ->
        {:error, {:unbalanced, %{program: unbalanced_program, missing_weight: correct_weight - unbalanced_weight}}}

      _ -> :ok
    end
  end

  defp read_programs_map(), do:
    File.stream!("input.txt")
    |> Stream.map(&String.trim/1)
    |> Stream.map(&decode_program/1)
    |> Stream.map(&{&1.name, &1})
    |> Enum.into(%{})

  defp decode_program(input) do
    [attributes | children] = String.split(input, " -> ")
    [name, weight] = String.split(attributes, " ")
    %{
      name: name,
      weight: weight |> String.replace(~r/(\(|\))/, "") |> String.to_integer(),
      children: decode_children(children)
    }
  end

  defp decode_children([]), do: []
  defp decode_children([children]), do:
    children
    |> String.replace(~r/\s/, "")
    |> String.split(",")
end

Day7.part1() |> IO.puts
Day7.part2() |> IO.inspect