{-# 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, 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