defmodule AOC.Day11 do
  use AOC.Day, day: 11

  def parse_input(lines) do
    Enum.map(lines, &String.to_charlist/1)
  end

  def part1(image), do: calculate_distance_sum(image, 2)

  def part2(image), do: calculate_distance_sum(image, 1_000_000)

  def calculate_distance_sum(image, scaler) do
    empty_space = find_empty_space(image)

    find_galaxy_positions(image)
    |> expand_galaxy_positions(empty_space, scaler)
    |> calculate_shortest_paths()
    |> Enum.sum()
  end

  def find_empty_space(image) do
    empty_rows = find_empty_rows(image)

    empty_columns =
      image
      |> List.zip()
      |> Enum.map(&Tuple.to_list/1)
      |> find_empty_rows()

    {empty_rows, empty_columns}
  end

  def find_empty_rows(image) do
    image
    |> Enum.with_index()
    |> Enum.filter(fn {row, _} -> Enum.all?(row, &(&1 == ?.)) end)
    |> Enum.map(&elem(&1, 1))
  end

  def find_galaxy_positions(image) do
    image
    |> Enum.map(fn row ->
      row
      |> Enum.with_index()
      |> Enum.filter(fn {space, _} -> space == ?# end)
      |> Enum.map(&elem(&1, 1))
    end)
    |> Enum.with_index()
    |> Enum.flat_map(fn {xs, y} ->
      Enum.map(xs, &{y, &1})
    end)
  end

  def expand_galaxy_positions(galaxy_positions, {empty_rows, empty_columns}, scaler) do
    Enum.map(galaxy_positions, fn {y, x} ->
      shift_y = Enum.count(empty_rows, &(&1 < y))
      shift_x = Enum.count(empty_columns, &(&1 < x))

      {y + shift_y * (scaler - 1), x + shift_x * (scaler - 1)}
    end)
  end

  def calculate_shortest_paths(galaxy_positions) do
    for g1 <- galaxy_positions, g2 <- galaxy_positions do
      [g1, g2]
    end
    |> Enum.reject(fn [g1, g2] -> g1 == g2 end)
    |> Enum.map(&Enum.sort/1)
    |> Enum.uniq()
    |> Enum.map(fn [{y1, x1}, {y2, x2}] ->
      abs(y1 - y2) + abs(x1 - x2)
    end)
  end
end