binaryfile.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module GeoPred.BinaryFile
   ( createOrLoadBinaryFile
   )
where

import Prelude hiding ((!))
import qualified Data.ByteString.Lazy.Char8 as BS
import qualified Data.Binary     as B
import qualified Data.Binary.Put as B
import qualified Data.Binary.Get as B
import System.Directory
import System.FilePath.Posix
import Data.List.Split
-- import System.IO
import Data.Array.Accelerate (Z(..), (:.)(..), fromList, Array(..))
import Data.Array.Accelerate.IO.Foreign.Ptr
import qualified Data.Vector.Generic.Mutable as VU hiding (length)
import qualified Data.Vector.Storable as VU
import Data.Vector.Storable ((!))
import Control.Monad
import Control.Monad.Trans.Class

readCsv :: FilePath -> IO [[Float]]
readCsv path =
  do contents <- BS.readFile (path ++ ".csv")
     let wrds = map (BS.splitWith (== ',')) . BS.lines $ contents
     let values :: [[Float]]
         values = map (map (read . BS.unpack) . tail ) $ wrds
     return values

createBinaryFloatFileFromCsv :: FilePath -> IO ()
createBinaryFloatFileFromCsv path =
  do putStrLn "Creating binary tracking data from .csv please wait..."
     values <- readCsv path
     putStrLn "writing entire tracking state file"
     let bs = writeFloats $ VU.fromList $ concat $ values
     BS.writeFile (path++".f32") bs
     let x = length values
         len = length (head values)
         z = 3
         y = x `div` z
         dataSetShape = Z :. x :. y :. z
     putStrLn "writing size"
     BS.writeFile (path++".size") . BS.pack . show $ [x, y, z]
     putStrLn "done"

writeFloats :: VU.Vector Float -> BS.ByteString
writeFloats v =
   B.runPut $ go 0
   where
   len :: Int
   len = VU.length v
   go :: Int -> B.Put
   go i =
       if i < len
       then do B.put (v ! i)
               go (i+1)
       else do B.flush

loadFloats :: Int -> BS.ByteString -> VU.Vector Float
loadFloats len bs =
  let go :: VU.Vector Float -> Int -> B.Get ()
      go v i =
        do empty <- B.isEmpty
           if empty
           then return ()
           else do a :: Float <- B.get
                   VU.write v i a
  in VU.createT $
          do vec <- VU.new len
             lift $ B.runGet (go vec 0) bs

loadBinaryFloatFile :: FilePath -> IO (Array (Z :. Int :. Int :. Int) Float)
loadBinaryFloatFile path =
  do putStrLn "Loading data from binary file..."
     [x, y, z] :: [Int] <- read . BS.unpack <$> BS.readFile (path ++ ".size")
     let sh = Z :. x :. y :. z
         len = x * y * z
     putStrLn $ "shape " ++ show sh
     bs <- BS.readFile (path++".f32")
     let values = loadFloats len bs
     arr <- VU.unsafeWith values (\ ptr -> return $ fromPtrs sh ptr)
     putStrLn "done"
     return arr

createOrLoadBinaryFile :: FilePath -> IO (Array (Z :. Int :. Int :. Int) Float)
createOrLoadBinaryFile pathRoot =
  do home <- getHomeDirectory
     let absolutePathRoot = normalise $ (++) home pathRoot
     exists <- doesFileExist (absolutePathRoot ++ ".f32")
     putStrLn absolutePathRoot
     when (not exists) $ createBinaryFloatFileFromCsv absolutePathRoot
     loadBinaryFloatFile absolutePathRoot