-- |
-- Module    : Statistics.Matrix.Types
-- Copyright : 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix.Types
    (
      Vector
    , MVector
    , Matrix(..)
    , MMatrix(..)
    , debug
    ) where

import Data.Char (isSpace)
import Numeric (showFFloat)
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as M

type Vector = U.Vector Double
type MVector s = M.MVector s Double

-- | Two-dimensional matrix, stored in row-major order.
data Matrix = Matrix {
      Matrix -> Int
rows     :: {-# UNPACK #-} !Int -- ^ Rows of matrix.
    , Matrix -> Int
cols     :: {-# UNPACK #-} !Int -- ^ Columns of matrix.
    , Matrix -> Vector
_vector  :: !Vector             -- ^ Matrix data.
    } deriving (Matrix -> Matrix -> Bool
(Matrix -> Matrix -> Bool)
-> (Matrix -> Matrix -> Bool) -> Eq Matrix
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Matrix -> Matrix -> Bool
$c/= :: Matrix -> Matrix -> Bool
== :: Matrix -> Matrix -> Bool
$c== :: Matrix -> Matrix -> Bool
Eq)

-- | Two-dimensional mutable matrix, stored in row-major order.
data MMatrix s = MMatrix
                 {-# UNPACK #-} !Int
                 {-# UNPACK #-} !Int
                 !(MVector s)

-- The Show instance is useful only for debugging.
instance Show Matrix where
    show :: Matrix -> String
show = Matrix -> String
debug

debug :: Matrix -> String
debug :: Matrix -> String
debug (Matrix Int
r Int
c Vector
vs) = [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (String -> ShowS) -> [String] -> [String] -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> ShowS
forall a. [a] -> [a] -> [a]
(++) (String
hdr0 String -> [String] -> [String]
forall a. a -> [a] -> [a]
: String -> [String]
forall a. a -> [a]
repeat String
hdr) [String]
rrows
  where
    rrows :: [String]
rrows         = ([String] -> String) -> [[String]] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (ShowS
cleanEnd ShowS -> ([String] -> String) -> [String] -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unwords) ([[String]] -> [String])
-> ([String] -> [[String]]) -> [String] -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> [[String]]
forall a. [a] -> [[a]]
split ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ (String -> ShowS) -> [String] -> [String] -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> ShowS
forall a. [a] -> [a] -> [a]
(++) [String]
ldone [String]
tdone
    hdr0 :: String
hdr0          = (Int, Int) -> String
forall a. Show a => a -> String
show (Int
r,Int
c) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" "
    hdr :: String
hdr           = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
hdr0) Char
' '
    pad :: (String -> t a -> t) -> Int -> t a -> t
pad String -> t a -> t
plus Int
k t a
xs = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
xs) Char
' ' String -> t a -> t
`plus` t a
xs
    ldone :: [String]
ldone         = ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> ShowS) -> Int -> ShowS
forall (t :: * -> *) a t.
Foldable t =>
(String -> t a -> t) -> Int -> t a -> t
pad String -> ShowS
forall a. [a] -> [a] -> [a]
(++) ([String] -> Int
forall a. [[a]] -> Int
longest [String]
lstr)) [String]
lstr
    tdone :: [String]
tdone         = ShowS -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((String -> ShowS) -> Int -> ShowS
forall (t :: * -> *) a t.
Foldable t =>
(String -> t a -> t) -> Int -> t a -> t
pad ((String -> ShowS) -> String -> ShowS
forall a b c. (a -> b -> c) -> b -> a -> c
flip String -> ShowS
forall a. [a] -> [a] -> [a]
(++)) ([String] -> Int
forall a. [[a]] -> Int
longest [String]
tstr)) [String]
tstr
    ([String]
lstr, [String]
tstr)  = [(String, String)] -> ([String], [String])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(String, String)] -> ([String], [String]))
-> (Vector -> [(String, String)]) -> Vector -> ([String], [String])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> (String, String)) -> [Double] -> [(String, String)]
forall a b. (a -> b) -> [a] -> [b]
map ((Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==Char
'.') (String -> (String, String))
-> (Double -> String) -> Double -> (String, String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> String
forall a. RealFloat a => a -> String
render) ([Double] -> [(String, String)])
-> (Vector -> [Double]) -> Vector -> [(String, String)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector -> [Double]
forall a. Unbox a => Vector a -> [a]
U.toList (Vector -> ([String], [String])) -> Vector -> ([String], [String])
forall a b. (a -> b) -> a -> b
$ Vector
vs
    longest :: [[a]] -> Int
longest       = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> ([[a]] -> [Int]) -> [[a]] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> Int) -> [[a]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
    render :: a -> String
render a
k      = ShowS
forall a. [a] -> [a]
reverse ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==Char
'.') ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==Char
'0') ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
forall a. [a] -> [a]
reverse ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                    Maybe Int -> a -> ShowS
forall a. RealFloat a => Maybe Int -> a -> ShowS
showFFloat (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
4) a
k ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String
""
    split :: [a] -> [[a]]
split []      = []
    split [a]
xs      = [a]
i [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
split [a]
rest where ([a]
i, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
c [a]
xs
    cleanEnd :: ShowS
cleanEnd      = ShowS
forall a. [a] -> [a]
reverse ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
forall a. [a] -> [a]
reverse