-module(almost_magic_square).
-export([almost_magic_square/1]).

almost_magic_square(A) ->
    Rows = chunk3(A),
    RowSums = [lists:sum(R) || R <- Rows],
    ColSums = [lists:sum([lists:nth(C, R) || R <- Rows]) || C <- [1, 2, 3]],
    MaxSum = lists:max(RowSums ++ ColSums),
    FinalRows = balance(0, 0, RowSums, ColSums, Rows, MaxSum),
    lists:append(FinalRows).

chunk3([]) -> [];
chunk3(List) ->
    {Row, Rest} = lists:split(3, List),
    [Row | chunk3(Rest)].

balance(I, J, RowSums, ColSums, Rows, MaxSum) when I < 3, J < 3 ->
    RowSumI = lists:nth(I + 1, RowSums),
    ColSumJ = lists:nth(J + 1, ColSums),
    Diff = min(MaxSum - RowSumI, MaxSum - ColSumJ),
    Rows1 = add_at(Rows, I, J, Diff),
    RowSums1 = add_at_list(RowSums, I, Diff),
    ColSums1 = add_at_list(ColSums, J, Diff),
    NextI = case lists:nth(I + 1, RowSums1) of MaxSum -> I + 1; _ -> I end,
    NextJ = case lists:nth(J + 1, ColSums1) of MaxSum -> J + 1; _ -> J end,
    balance(NextI, NextJ, RowSums1, ColSums1, Rows1, MaxSum);
balance(_, _, _, _, Rows, _) ->
    Rows.

add_at(Rows, I, J, Diff) ->
    [case Idx of
         I -> [case Jdx of J -> V + Diff; _ -> V end
               || {Jdx, V} <- lists:zip(lists:seq(0, length(Row) - 1), Row)];
         _ -> Row
     end || {Idx, Row} <- lists:zip(lists:seq(0, length(Rows) - 1), Rows)].

add_at_list(List, Idx0, Diff) ->
    [case I of Idx0 -> V + Diff; _ -> V end
     || {I, V} <- lists:zip(lists:seq(0, length(List) - 1), List)].