almostMagicSquare :: [Int] -> [Int]
almostMagicSquare a = concat (go rows rowSums colSums 0 0)
  where
    chunk3 [] = []
    chunk3 xs = take 3 xs : chunk3 (drop 3 xs)

    updateAt :: Int -> (b -> b) -> [b] -> [b]
    updateAt idx f xs = [if i == idx then f x else x | (i, x) <- zip [0 ..] xs]

    rows    = chunk3 a
    rowSums = map sum rows
    colSums = [sum [rows !! r !! c | r <- [0 .. 2]] | c <- [0 .. 2]]
    maxSum  = maximum (rowSums ++ colSums)

    go grid rs cs i j
      | i >= 3 || j >= 3 = grid
      | otherwise =
          let diff  = min (maxSum - rs !! i) (maxSum - cs !! j)
              grid' = updateAt i (updateAt j (+ diff)) grid
              rs'   = updateAt i (+ diff) rs
              cs'   = updateAt j (+ diff) cs
              i'    = if rs' !! i == maxSum then i + 1 else i
              j'    = if cs' !! j == maxSum then j + 1 else j
          in  go grid' rs' cs' i' j'