content: {-# LANGUAGE FlexibleContexts #-}
  
  module AI.Util.Matrix where
  
  import Control.Monad.Random hiding (fromList)
  import Foreign.Storable (Storable)
  import Numeric.LinearAlgebra
  
  -- |Return the size of a matrix as a 2-tuple.
  size :: Matrix a -> (Int,Int)
  size x = (rows x, cols x)
  
  -- |Concatenate matrices horizontally.
  horzcat :: Element a => [Matrix a] -> Matrix a
  horzcat = fromBlocks . return
  
  -- |Concatenate matrices vertically.
  vertcat :: Element a => [Matrix a] -> Matrix a
  vertcat = fromBlocks . map return
  
  -- |Add a column of ones to a matrix.
  addOnes :: Matrix Double -> Matrix Double
  addOnes x = fromBlocks [[1, x]]
  
  -- |Create a row matrix.
  row :: [Double] -> Matrix Double
  row = asRow . fromList
  
  -- |Create a column matrix.
  column :: [Double] -> Matrix Double
  column = asColumn . fromList
  
  --------------------------
  -- Functions on Vectors --
  --------------------------
  
  takeVector :: Storable a => Int -> Vector a -> Vector a
  takeVector n v = subVector 0 n v
  
  dropVector :: Storable a => Int -> Vector a -> Vector a
  dropVector n v = subVector n (dim v - n) v
  
  sumVector :: (Num a, Storable a) => Vector a -> a
  sumVector xs = foldVector (+) 0 xs
  
  prodVector :: (Num a, Storable a) => Vector a -> a
  prodVector xs = foldVector (*) 1 xs
  
  ---------------------------
  -- Functions on Matrices --
  ---------------------------
  
  mapRows :: Element a => (Vector a -> b) -> Matrix a -> [b]
  mapRows f m = map f (toRows m)
  
  mapCols :: Element a => (Vector a -> b) -> Matrix a -> [b]
  mapCols f m = map f (toColumns m)
  
  eachRow :: (Element a, Element b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
  eachRow f = fromRows . mapRows f
  
  eachCol :: (Element a, Element b) => (Vector a -> Vector b) -> Matrix a -> Matrix b
  eachCol f = fromColumns . mapCols f
  
  sumRows :: (Element a, Num (Vector a)) => Matrix a -> Vector a
  sumRows m = sum $ toRows m
  
  sumCols :: (Element a, Num (Vector a)) => Matrix a -> Vector a
  sumCols m = sum $ toColumns m
  
  sumMatrix :: Matrix Double -> Double
  sumMatrix = sumVector . sum . toRows
  
  ------------------------
  -- Subset Referencing --
  ------------------------
  
  subRefVec :: Storable a => Vector a -> [Int] -> Vector a
  subRefVec v is = fromList $ map (v@>) is
  
  subRefRows :: Element a => Matrix a -> [Int] -> Matrix a
  subRefRows m is = fromRows $ map (r!!) is where r = toRows m
  
  subRefCols :: Element a => Matrix a -> [Int] -> Matrix a
  subRefCols m is = fromColumns $ map (c!!) is where c = toColumns m
  
  