content: {-# LANGUAGE CPP, BangPatterns #-}
  {-# OPTIONS_GHC -Wall -fno-warn-name-shadowing #-}
  
  module Main ( main, test {-, maxDistances -} ) where
  
  import Prelude
  import System.Environment
  import Data.Array.Accelerate as A
  import Data.Array.Accelerate.CUDA
  import AccelerateCompat as A
  
  -- <<Graph
  type Weight = Int32
  type Graph = Array DIM2 Weight
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- shortestPaths
  
  -- <<shortestPaths
  shortestPaths :: Graph -> Graph
  shortestPaths g0 = run (shortestPathsAcc n (use g0))
    where
      Z :. _ :. n = arrayShape g0
  -- >>
  
  -- <<shortestPathsAcc
  shortestPathsAcc :: Int -> Acc Graph -> Acc Graph
  shortestPathsAcc n g0 = foldl1 (>->) steps g0              -- <3>
   where
    steps :: [ Acc Graph -> Acc Graph ]                      -- <1>
    steps =  [ step (unit (constant k)) | k <- [0 .. n-1] ]  -- <2>
  -- >>
  
  -- <<step
  step :: Acc (Scalar Int) -> Acc Graph -> Acc Graph
  step k g = generate (shape g) sp                           -- <1>
   where
     k' = the k                                              -- <2>
  
     sp :: Exp DIM2 -> Exp Weight
     sp ix = let
               (Z :. i :. j) = unlift ix                     -- <3>
             in
               A.min (g ! (index2 i j))                      -- <4>
                     (g ! (index2 i k') + g ! (index2 k' j))
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- Testing
  
  -- <<inf
  inf :: Weight
  inf = 999
  -- >>
  
  testGraph :: Graph
  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
  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 = toList (shortestPaths testGraph) == toList expectedResult
  
  toAdjMatrix :: [[Weight]] -> Graph
  toAdjMatrix xs = A.fromList (Z :. k :. k) (concat xs)
    where k = length xs
  
  
  main :: IO ()
  main = do
     (n:_) <- fmap (fmap read) getArgs
     print (run (let g :: Acc Graph
                     g    = generate (constant (Z:.n:.n) :: Exp DIM2) f
  
                     f :: Exp DIM2 -> Exp Weight
                     f ix = let i,j :: Exp Int
                                Z:.i:.j = unlift ix
                            in
                            A.fromIntegral j +
                             A.fromIntegral i * constant (Prelude.fromIntegral n)
                 in
                 A.foldAll (+) (constant 0) (shortestPathsAcc n g)))
  