content: {-# LANGUAGE CPP, BangPatterns #-}
  {-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
  
  module Main ( main, test ) where
  
  import System.Environment
  import Data.Array.Repa
  import Data.Functor.Identity
  
  -- <<Graph
  type Weight = Int
  type Graph r = Array r DIM2 Weight
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- shortestPaths
  
  -- <<shortestPaths
  shortestPaths :: Graph U -> Graph U
  shortestPaths g0 = runIdentity $ go g0 0                      -- <1>
    where
      Z :. _ :. n = extent g0
  
      go !g !k | k == n    = return g                           -- <2>
               | otherwise = do
                   g' <- computeP (fromFunction (Z:.n:.n) sp)   -- <3>
                   go g' (k+1)
       where
          sp (Z:.i:.j) = min (g ! (Z:.i:.j)) (g ! (Z:.i:.k) + g ! (Z:.k:.j))
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- Testing
  
  input :: [[Int]]
  input = [[  0, 999, 999,  13, 999, 999],
           [999,   0, 999, 999,   4,   9],
           [ 11, 999,   0, 999, 999, 999],
           [999,   3, 999,   0, 999,   7],
           [ 15,   5, 999,   1,   0, 999],
           [ 11, 999, 999,  14, 999,   0]]
  
  -- correct result:
  result :: [[Int]]
  result = [[0,  16, 999, 13, 20, 20],
            [19,  0, 999,  5,  4,  9],
            [11, 27,   0, 24, 31, 31],
            [18,  3, 999,  0,  7,  7],
            [15,  4, 999,  1,  0,  8],
            [11, 17, 999, 14, 21,  0] ]
  
  test :: Bool
  test = fromAdjMatrix (shortestPaths (toAdjMatrix input)) == result
  
  toAdjMatrix :: [[Int]] -> Graph U
  toAdjMatrix xs = fromListUnboxed (Z :. k :. k) (concat xs)
    where k = length xs
  
  fromAdjMatrix :: Graph U -> [[Int]]
  fromAdjMatrix m = chunk k (toList m)
    where
     (Z :. _ :. k) = extent m
  
  chunk :: Int -> [a] -> [[a]]
  chunk _ [] = []
  chunk n xs = as : chunk n bs
    where (as,bs) = splitAt n xs
  
  main :: IO ()
  main = do
     [n] <- fmap (fmap read) getArgs
     let g = fromListUnboxed (Z:.n:.n) [1..n^(2::Int)] :: Graph U
     print (sumAllS (shortestPaths g))
  
  