{-# LANGUAGE PackageImports, FlexibleContexts #-}
-- | Read and write matrices as ASCII text files.
--
--   The file format is like:
--
--   @
--      MATRIX                  -- header
--      100 100                 -- width and height
--      1.23 1.56 1.23 ...      -- data, separated by whitespace
--      ....
--   @
module Data.Array.Repa.IO.Matrix
        ( readMatrixFromTextFile
        , writeMatrixToTextFile)
where
import Data.Array.Repa.IO.Internals.Text
import Control.Monad
import System.IO
import Data.List                                as L
import Data.Array.Repa                          as A
import Data.Array.Repa.Repr.Unboxed             as A
import Prelude                                  as P


-- | Read a matrix from a text file.
-- 
--   * WARNING: This is implemented fairly naively, just using `Strings` 
--     under the covers. It will be slow for large data files.
-- 
--   * It also doesn't do graceful error handling.
--     If the file has the wrong format you'll get a confusing `error`.
--
readMatrixFromTextFile
        :: (Num e, Read e, Unbox e)
        => FilePath
        -> IO (Array U DIM2 e)  

readMatrixFromTextFile :: forall e.
(Num e, Read e, Unbox e) =>
FilePath -> IO (Array U DIM2 e)
readMatrixFromTextFile FilePath
fileName
 = do   Handle
handle          <- FilePath -> IOMode -> IO Handle
openFile FilePath
fileName IOMode
ReadMode
        
        FilePath
"MATRIX"        <- Handle -> IO FilePath
hGetLine Handle
handle
        [Int
width, Int
height] <- (FilePath -> [Int]) -> IO FilePath -> IO [Int]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ((FilePath -> Int) -> [FilePath] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
P.map FilePath -> Int
forall a. Read a => FilePath -> a
read ([FilePath] -> [Int])
-> (FilePath -> [FilePath]) -> FilePath -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> [FilePath]
words) (IO FilePath -> IO [Int]) -> IO FilePath -> IO [Int]
forall a b. (a -> b) -> a -> b
$ Handle -> IO FilePath
hGetLine Handle
handle
        FilePath
str             <- Handle -> IO FilePath
hGetContents Handle
handle
        let vals :: [e]
vals        = FilePath -> [e]
forall a. (Num a, Read a) => FilePath -> [a]
readValues FilePath
str

        let dims :: DIM2
dims        = DIM0
Z DIM0 -> Int -> DIM0 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
width (DIM0 :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
height
        Array U DIM2 e -> IO (Array U DIM2 e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array U DIM2 e -> IO (Array U DIM2 e))
-> Array U DIM2 e -> IO (Array U DIM2 e)
forall a b. (a -> b) -> a -> b
$ DIM2 -> [e] -> Array U DIM2 e
forall sh a. (Shape sh, Unbox a) => sh -> [a] -> Array U sh a
fromListUnboxed DIM2
dims [e]
vals


-- | Write a matrix as a text file.
writeMatrixToTextFile 
        :: (Show e, Source r e)
        => FilePath
        -> Array r DIM2 e
        -> IO ()

writeMatrixToTextFile :: forall e r.
(Show e, Source r e) =>
FilePath -> Array r DIM2 e -> IO ()
writeMatrixToTextFile FilePath
fileName Array r DIM2 e
arr
 = do   Handle
file    <- FilePath -> IOMode -> IO Handle
openFile FilePath
fileName IOMode
WriteMode  

        Handle -> FilePath -> IO ()
hPutStrLn Handle
file FilePath
"MATRIX"

        let DIM0
Z :. Int
width :. Int
height        
                = Array r DIM2 e -> DIM2
forall sh. Shape sh => Array r sh e -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array r DIM2 e
arr

        Handle -> FilePath -> IO ()
hPutStrLn Handle
file (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> FilePath
forall a. Show a => a -> FilePath
show Int
width FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
P.++ FilePath
" " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
P.++ Int -> FilePath
forall a. Show a => a -> FilePath
show Int
height
        Handle -> [e] -> IO ()
forall a. Show a => Handle -> [a] -> IO ()
hWriteValues Handle
file ([e] -> IO ()) -> [e] -> IO ()
forall a b. (a -> b) -> a -> b
$ Array r DIM2 e -> [e]
forall sh r e. (Shape sh, Source r e) => Array r sh e -> [e]
toList Array r DIM2 e
arr
        Handle -> IO ()
hClose Handle
file