content: {-# LANGUAGE FlexibleContexts #-}
  
  module SMVM where
  
  import Random
  import SMVM.Matrix
  
  import System.Random.MWC
  import Data.Array.Unboxed
  import Data.Array.Accelerate           (Vector, Segments, Acc)
  import qualified Data.Array.Accelerate as Acc
  import qualified Data.Vector.Unboxed   as V
  
  
  -- Sparse-matrix vector multiplication
  -- -----------------------------------
  
  type SparseVector a = (Vector Int, Vector a)
  type SparseMatrix a = (Segments Int, SparseVector a)
  
  smvmAcc :: SparseMatrix Float -> Vector Float -> Acc (Vector Float)
  smvmAcc (segd', (inds', vals')) vec'
    = let
        segd     = Acc.use segd'
        inds     = Acc.use inds'
        vals     = Acc.use vals'
        vec      = Acc.use vec'
        ---
        vecVals  = Acc.backpermute (Acc.shape inds) (\i -> Acc.index1 $ inds Acc.! i) vec
        products = Acc.zipWith (*) vecVals vals
      in
      Acc.foldSeg (+) 0 products segd
  
  
  -- The reference version will be slow, with many conversions between
  -- array/vector/list representations. This will likely skew heap usage
  -- calculations, but oh well...
  --
  type USparseMatrix a = (UArray Int Int, (UArray Int Int, UArray Int a))
  
  smvmRef :: USparseMatrix Float -> UArray Int Float -> UArray Int Float
  smvmRef (segd, (inds, values)) vec
    = listArray (0, rangeSize (bounds segd) - 1)
      [sum [ values!i * vec!(inds!i) | i <- range seg] | seg <- segd' ]
    where
      segbegin = scanl  (+) 0 $ elems segd
      segend   = scanl1 (+)   $ elems segd
      segd'    = zipWith (\x y -> (x,y-1)) segbegin segend
  
  
  -- Main
  -- ----
  
  run :: Maybe FilePath -> IO (() -> UArray Int Float, () -> Acc (Vector Float))
  run f = withSystemRandom $ \gen ->  do
    -- sparse-matrix
    (segd', smat') <- maybe (randomCSRMatrix gen 512 512) (readCSRMatrix gen) f
    let (ind',val') = V.unzip smat'
  
    segd <- convertVector segd'
    ind  <- convertVector ind'
    val  <- convertVector val'
    let smat = (segd, (ind,val))
  
    -- vector
    vec' <- uniformVector gen (V.length segd')
    vec  <- convertVector vec'
  
    -- multiply!
    return (run_ref (v2a segd', (v2a ind',v2a val')) (v2a vec'), run_acc smat vec)
    where
      {-# NOINLINE run_ref #-}
      run_ref smat vec () = smvmRef smat vec
      run_acc smat vec () = smvmAcc smat vec
      --
      v2a :: (V.Unbox a, IArray UArray a) => V.Vector a -> UArray Int a
      v2a vec = listArray (0, V.length vec - 1) $ V.toList vec
  
  