{-# LANGUAGE BangPatterns, TupleSections #-}

module SMVM.Matrix (readCSRMatrix, randomCSRMatrix) where

import Random
import SMVM.MatrixMarket
import System.Random.MWC
import System.IO.Unsafe

import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed            as V
import qualified Data.Vector.Unboxed.Mutable    as M
import qualified Data.Vector.Algorithms.Intro   as V

type CSRMatrix a = (Vector Int, Vector (Int,a))


-- Read a sparse matrix from a MatrixMarket file. Pattern matrices are filled
-- with random numbers in the range (-1,1).
--
{-# INLINE readCSRMatrix #-}
readCSRMatrix :: GenIO -> FilePath -> IO (CSRMatrix Float)
readCSRMatrix gen file = do
  mtx <- readMatrix file
  case mtx of
    (RealMatrix    dim l vals) -> csr dim l vals
    (PatternMatrix dim l ix)   -> csr dim l =<< mapM' (\(a,b) -> (a,b,) `fmap` uniformR (-1,1) gen) ix
    (IntMatrix _ _ _)          -> error "IntMatrix type not supported"
    (ComplexMatrix _ _ _)      -> error "ComplexMatrix type not supported"


-- A randomly generated matrix of given size
--
{-# INLINE randomCSRMatrix #-}
randomCSRMatrix :: GenIO -> Int -> Int -> IO (CSRMatrix Float)
randomCSRMatrix gen rows cols = do
  segd <- randomVectorR ( 0,cols`div`3) gen rows
  let nnz = V.sum segd
  inds <- randomVectorR ( 0,cols-1) gen nnz
  vals <- randomVectorR (-1,1) gen nnz
  return (segd, V.zip inds vals)


-- Read elements into unboxed arrays, convert to zero-indexed compress sparse
-- row format.
--
{-# INLINE csr #-}
csr :: (Int,Int) -> Int -> [(Int,Int,Float)] -> IO (Vector Int, Vector (Int,Float))
csr (m,_) l elems = do
  mu <- M.new l
  let goe  _ []     = return ()
      goe !n (x:xs) = let (i,j,v) = x in M.unsafeWrite mu n (i-1,j-1,v) >> goe (n+1) xs
  goe 0 elems

  let cmp (x1,y1,_) (x2,y2,_) | x1 == x2  = compare y1 y2
                              | otherwise = compare x1 x2
  V.sortBy cmp mu

  (i,j,v) <- V.unzip3 `fmap` V.unsafeFreeze mu
  mseg    <- M.new m
  let gos !n rows | n < m     = let (s,ss) = V.span (==n) rows in M.unsafeWrite mseg n (V.length s) >> gos (n+1) ss
                  | otherwise = V.unsafeFreeze mseg
  seg <- gos 0 i
  return (seg , V.zip j v)


-- Lazier versions of things in Control.Monad
--
sequence' :: [IO a] -> IO [a]
sequence' ms = foldr k (return []) ms
    where k m m' = do { x <- m; xs <- unsafeInterleaveIO m'; return (x:xs) }

mapM' :: (a -> IO b) -> [a] -> IO [b]
mapM' f as = sequence' (map f as)