module Main (main) where

import qualified Data.StorableVector.Stream as VG
import qualified Data.StorableVector as V

import qualified Data.StorableVector.Lazy.Stream as VGL
import qualified Data.StorableVector.Lazy as VL

import qualified Data.ByteString.Lazy as B
import qualified Data.Binary.Builder as Bin
import qualified Data.Monoid as Md

import qualified Data.List.Stream as L
import qualified Data.Stream as S
import Control.Monad.Stream (zipWithM_, )

import Control.Monad (when, )
import System.IO (withBinaryFile, hPutBuf, Handle, IOMode(WriteMode))
import Foreign (Storable, Ptr, Int16, poke, allocaArray, sizeOf, advancePtr, )

import GHC.Float (double2Int, )

import System.Time (getClockTime, diffClockTimes, tdSec, tdPicosec, )



{-# INLINE signalToBinaryPut #-}
signalToBinaryPut :: [Int16] -> B.ByteString
signalToBinaryPut =
   Bin.toLazyByteString . mconcat .
   L.map (Bin.putWord16host . fromIntegral)

mconcat :: Md.Monoid m => [m] -> m
mconcat =
   L.foldr Md.mappend Md.mempty


{- awfully slow
doubleToInt16 :: Double -> Int16
doubleToInt16 x = round (32767 * x)
-}

round' :: Double -> Int16
round' x =
   fromIntegral (double2Int
     (if x<0 then x-0.5 else x+0.5))

doubleToInt16 :: Double -> Int16
doubleToInt16 x = round' (32767 * x)


-- that's important in order to allow fusion
{-# INLINE exponential2 #-}
exponential2 :: Double -> Double -> [Double]
exponential2 hl y0 =
   let k = 0.5 ** recip hl
   in  L.iterate (k*) y0

exponential2Gen :: Double -> Double -> S.Stream Double
exponential2Gen hl y0 =
   let k = 0.5 ** recip hl
   in  S.iterate (k*) y0


{-# INLINE writeSignal #-}
writeSignal :: FilePath -> Int -> [Double] -> IO ()
writeSignal name num signal =
   withBinaryFile name WriteMode $ \h ->
   allocaArray num $ \buf ->
      zipWithM_ poke
         (L.take num $ L.iterate (flip advancePtr 1) buf)
         (L.map doubleToInt16 signal) >>
      hPutArray h buf num

writeExponentialList :: FilePath -> Int -> Double -> Double -> IO ()
writeExponentialList name num hl y0 =
   withBinaryFile name WriteMode $ \h ->
   allocaArray num $ \buf ->
      zipWithM_ poke
         (L.take num $ L.iterate (flip advancePtr 1) buf)
         (L.map doubleToInt16
             (let k = 0.5 ** recip hl
              in  L.iterate (k*) y0)) >>
      hPutArray h buf num

writeExponential :: FilePath -> Int -> Double -> Double -> IO ()
writeExponential name num hl y0 =
   withBinaryFile name WriteMode $ \h ->
   allocaArray num $ \buf ->
      let k = 0.5 ** recip hl
          endPtr = advancePtr buf num
          loop ptr y =
             when (ptr<endPtr) $
                poke ptr (doubleToInt16 y) >>
                loop (advancePtr ptr 1) (y*k)
      in  loop buf y0 >>
          hPutArray h buf num

hPutArray :: Storable a => Handle -> Ptr a -> Int -> IO ()
hPutArray h buf num =
   let size :: Storable a => Ptr a -> a -> Int
       size _ dummy = num * sizeOf dummy
   in  hPutBuf h buf (size buf undefined)

exponentialStorableVector :: Int -> Double -> Double -> V.Vector Int16
exponentialStorableVector num hl y0 =
   let k = 0.5 ** recip hl
   in  fst $ V.unfoldrN num (\y -> Just (doubleToInt16 y, y*k)) y0

exponentialStorableVectorLazy :: Double -> Double -> VL.Vector Int16
exponentialStorableVectorLazy hl y0 =
   let k = 0.5 ** recip hl
   in  VL.unfoldr VL.defaultChunkSize
          (\y -> Just (doubleToInt16 y, y*k)) y0



measureTime :: String -> (FilePath -> IO ()) -> IO ()
measureTime name act =
   do putStr (name++": ")
      timeA <- getClockTime
      act (name++".sw")
      timeB <- getClockTime
      let td = diffClockTimes timeB timeA
      print (fromIntegral (tdSec td) +
             fromInteger (tdPicosec td) * 1e-12 :: Double)

numSamples :: Int
numSamples = 1000000

halfLife :: Double
halfLife = 100000


main :: IO ()
main =
   do measureTime "storablevector-from-stream" $ \fn ->
         V.writeFile fn $
         VG.from numSamples $
         S.map doubleToInt16 $
         exponential2Gen halfLife 1
      measureTime "storablevector-fused" $ \fn ->
         V.writeFile fn $
         VG.fromList numSamples $
         L.map doubleToInt16 $
         exponential2 halfLife 1
      measureTime "storablevector" $ \fn ->
         V.writeFile fn $
         exponentialStorableVector numSamples halfLife 1

      measureTime "storablevector-lazy-from-stream" $ \fn ->
         VL.writeFile fn $
         VGL.from VL.defaultChunkSize $
         S.take numSamples $
         S.map doubleToInt16 $
         exponential2Gen halfLife 1
      measureTime "storablevector-lazy-fused" $ \fn ->
         VL.writeFile fn $
         VGL.fromList VL.defaultChunkSize $
         L.take numSamples $
         L.map doubleToInt16 $
         exponential2 halfLife 1
      measureTime "storablevector-lazy" $ \fn ->
         VL.writeFile fn $
         VL.take numSamples $
         exponentialStorableVectorLazy halfLife 1

      measureTime "loop-poke" $ \fn ->
         writeExponential fn numSamples halfLife 1
      measureTime "list-binary-put" $ \fn ->
         B.writeFile fn $
         signalToBinaryPut $
         L.take numSamples $
         L.map doubleToInt16 $
         exponential2 halfLife 1
      measureTime "list-poke-buffer" $ \fn ->
         writeSignal fn numSamples $
         exponential2 halfLife 1
      measureTime "list-poke" $ \fn ->
         writeExponentialList fn numSamples halfLife 1