defmodule AlmostMagicSquare do
  def almost_magic_square(a) do
    rows = Enum.chunk_every(a, 3)
    row_sums = Enum.map(rows, &Enum.sum/1)
    col_sums = for j <- 0..2, do: rows |> Enum.map(&Enum.at(&1, j)) |> Enum.sum()
    max_sum = Enum.max(row_sums ++ col_sums)

    {final_rows, _, _, _, _} = balance(rows, row_sums, col_sums, max_sum, 0, 0)
    List.flatten(final_rows)
  end

  defp balance(rows, row_sums, col_sums, max_sum, i, j) when i < 3 and j < 3 do
    diff = min(max_sum - Enum.at(row_sums, i), max_sum - Enum.at(col_sums, j))

    rows = List.update_at(rows, i, fn row -> List.update_at(row, j, &(&1 + diff)) end)
    row_sums = List.update_at(row_sums, i, &(&1 + diff))
    col_sums = List.update_at(col_sums, j, &(&1 + diff))

    next_i = if Enum.at(row_sums, i) == max_sum, do: i + 1, else: i
    next_j = if Enum.at(col_sums, j) == max_sum, do: j + 1, else: j

    balance(rows, row_sums, col_sums, max_sum, next_i, next_j)
  end

  defp balance(rows, row_sums, col_sums, _max_sum, i, j), do: {rows, row_sums, col_sums, i, j}
end