{-# LANGUAGE FlexibleContexts #-}

-- | The Volumetric module contains functions for manipulating the
--   volumetric data files found at,
--
--     <http://graphics.stanford.edu/data/voldata/>
--
module Volumetric (
  bracket_array,
  flip_x,
  flip_y,
  read_word16s,
  round_array,
  swap_bytes,
  write_values_to_bmp,
  write_word16s,
  z_slice )
where

import Data.Word ( Word8, Word16 )
import Data.Bits ( (.&.), (.|.), shift )
import qualified Data.Array.Repa as Repa ( map )
import Data.Array.Repa (
  (:.)( (:.) ),
  Array,
  D,
  DIM2,
  DIM3,
  Shape,
  Source,
  U,
  Z( Z ),
  computeUnboxedP,
  copyP,
  slice )
import Data.Array.Repa.Algorithms.ColorRamp ( rampColorHotToCold )
import Data.Array.Repa.Eval ( now )
import Data.Array.Repa.IO.Binary (
  readArrayFromStorableFile,
  writeArrayToStorableFile )
import Data.Array.Repa.IO.BMP ( writeImageToBMP )
import Data.Array.Repa.Operators.Traversal ( unsafeTraverse )
import Data.Array.Repa.Slice ( All( All ), Any( Any ) )

import Values ( Values, Values2D )


-- | RawData is an array of words (16 bits), as contained in the
--   volumetric data files.
type RawData sh = Array U sh Word16

-- | A specialization of the 'RawData' type, to three dimensions.
type RawData3D = RawData DIM3

type RGB = (Word8, Word8, Word8)
type ColorData sh = Array U sh RGB


{-# INLINE read_word16s #-}
read_word16s :: FilePath -> DIM3 -> IO RawData3D
read_word16s path shape = do
  arr <- readArrayFromStorableFile path shape
  c   <- copyP arr
  now c



bracket :: Double -> Double -> Double -> Word16
bracket lower_threshold upper_threshold x
        | x < lower_threshold      = 0
        | x > upper_threshold      = 255
        | otherwise                = truncate (r * 255)
            where
              numerator = x - lower_threshold
              denominator = upper_threshold - lower_threshold
              r = numerator/denominator


flip16 :: Word16 -> Word16
flip16 xx =
  shift xx 8 .|. (shift xx (-8) .&. 0x00ff)


{-# INLINE swap_bytes #-}
swap_bytes :: (Shape sh, Source r Word16) => Array r sh Word16
             -> Array D sh Word16
swap_bytes =
  Repa.map flip16


bracket_array :: Shape sh => Double -> Double -> Values sh -> Array D sh Word16
bracket_array lt ut =
  Repa.map (bracket lt ut)


{-# INLINE round_array #-}
round_array :: Shape sh => Values sh -> Array D sh Word16
round_array =
  Repa.map round


flip_y :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
flip_y height arr =
  unsafeTraverse arr id
              (\get (Z :. z :. y :. x) ->
                   get (Z :. z :. (height - 1) - y :. x))

flip_x :: Source r Word16 => Int -> Array r DIM3 Word16 -> Array D DIM3 Word16
flip_x width arr =
  unsafeTraverse arr id
              (\get (Z :. z :. y :. x) ->
                   get (Z :. z :. y :. (width - 1) - x))


{-# INLINE write_word16s #-}
write_word16s :: (Shape sh) => FilePath -> (RawData sh) -> IO ()
write_word16s = writeArrayToStorableFile



--
-- Instead of IO, we could get away with a generic monad 'm'
-- here. However, /we/ only call this function from within IO.
--
values_to_colors :: (Shape sh) => (Values sh) -> IO (ColorData sh)
values_to_colors arr =
  computeUnboxedP $ Repa.map (truncate_rgb . ramp_it) arr
  where
    ramp_it :: Double -> (Double, Double, Double)
    ramp_it x =
      if x == 0 then
        (0, 0, 0)
      else
        rampColorHotToCold 0 255 x

    truncate_rgb :: (Double, Double, Double) -> (Word8, Word8, Word8)
    truncate_rgb (r, g, b) =
      (r', g', b')
      where
        r' = truncate (r * 255) :: Word8
        g' = truncate (g * 255) :: Word8
        b' = truncate (b * 255) :: Word8


write_values_to_bmp :: FilePath -> Values2D -> IO ()
write_values_to_bmp path values = do
  colors <- values_to_colors values
  writeImageToBMP path colors


z_slice :: (Source r a) => Int -> Array r DIM3 a -> Array D DIM2 a
z_slice n arr =
  slice arr (Any :. n :. All :. All)