content: {-# LANGUAGE CPP, BangPatterns #-}
  {-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
  
  module Main ( main, test, maxDistances ) where
  
  import System.Environment
  import Data.Array.Repa
  
  -- <<Graph
  type Weight = Int
  type Graph r = Array r DIM2 Weight
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- shortestPaths
  
  -- <<shortestPaths
  shortestPaths :: Graph U -> Graph U
  shortestPaths g0 = go g0 0                                                -- <2>
    where
      Z :. _ :. n = extent g0                                               -- <1>
  
      go !g !k | k == n    = g                                              -- <3>
               | otherwise =
                   let g' = computeS (fromFunction (Z:.n:.n) sp)            -- <4>
                   in  go g' (k+1)                                          -- <5>
       where
         sp (Z:.i:.j) = min (g ! (Z:.i:.j)) (g ! (Z:.i:.k) + g ! (Z:.k:.j)) -- <6>
  -- >>
  
  -- -----------------------------------------------------------------------------
  
  -- <<maxDistance
  maxDistance :: Weight -> Weight -> Weight
  maxDistance x y
    | x == inf  = y
    | y == inf  = x
    | otherwise = max x y
  -- >>
  
  maxDistances :: Graph U -> Array U DIM1 Weight
  maxDistances = foldS maxDistance inf
  
  -- -----------------------------------------------------------------------------
  -- Testing
  
  -- <<inf
  inf :: Weight
  inf = 999
  -- >>
  
  testGraph :: Graph U
  testGraph = toAdjMatrix $
          [[  0, inf, inf,  13, inf, inf],
           [inf,   0, inf, inf,   4,   9],
           [ 11, inf,   0, inf, inf, inf],
           [inf,   3, inf,   0, inf,   7],
           [ 15,   5, inf,   1,   0, inf],
           [ 11, inf, inf,  14, inf,   0]]
  
  -- correct result:
  expectedResult :: Graph U
  expectedResult = toAdjMatrix $
           [[0,  16, inf, 13, 20, 20],
            [19,  0, inf,  5,  4,  9],
            [11, 27,   0, 24, 31, 31],
            [18,  3, inf,  0,  7,  7],
            [15,  4, inf,  1,  0,  8],
            [11, 17, inf, 14, 21,  0] ]
  
  test :: Bool
  test = shortestPaths testGraph == expectedResult
  
  toAdjMatrix :: [[Weight]] -> Graph U
  toAdjMatrix xs = fromListUnboxed (Z :. k :. k) (concat xs)
    where k = length xs
  
  main :: IO ()
  main = do
     [n] <- fmap (fmap read) getArgs
     let g = fromListUnboxed (Z:.n:.n) [0..n^(2::Int)-1] :: Graph U
     print (sumAllS (shortestPaths g))
  
  