{-# LANGUAGE BangPatterns #-}
module Graphics.Image.PFS where

import Foreign.C
import Foreign
import Graphics.Image.PFS.Internal
import qualified Graphics.Image.PixelMap as P
import qualified Data.Array.Storable as A
import qualified Data.Map as M
import Control.Monad
import System.IO
import System.Posix.IO
import Control.Applicative

-- | Datatype representing a floating point image frame
data FPFrame = FPFrame {
    tags :: [(String,String)]
  , channelX :: A.StorableArray (Int,Int) CFloat
  , channelY :: A.StorableArray (Int,Int) CFloat
  , channelZ :: A.StorableArray (Int,Int) CFloat
  , channels :: M.Map String (A.StorableArray (Int,Int) CFloat)
  , width :: Int
  , height :: Int
}

-- | To the portable PixelMap format
instance P.ImageData FPFrame where
    toPixelMap (FPFrame t x y z cs w h) = P.PixelMap 
        { P.width=w
        , P.height=h
        , P.pixels=M.fromList $ [("X", x), ("Y",y),("Z",z)] ++ M.toList cs
        , P.colorspace=P.Ciea
        , P.tags=t }
    fromPixelMap (P.PixelMap t p w h P.Ciea) = FPFrame t (p M.! "X") (p M.! "Y") (p M.! "Z") M.empty w h        

-- | Read a frame from a handle
hReadFrame :: Handle -> IO FPFrame
hReadFrame fileH = do 
    fileFd <- fromIntegral <$> handleToFd fileH
    fileh <- withCString "rb" $ \filemode -> fdopen fileFd filemode
    domio <- pfs_newDOMIO
    
    frame <- pfs_DOMIOReadFrame domio fileh
    tagctr <- pfs_FrameGetTags frame
    tagitr <- pfs_TagContainerGetIterator tagctr
    let loop !k = if (pfs_TagIteratorHasNext tagitr)==0 
                    then do tagname <- pfs_TagIteratorGetNext tagitr 
                            tagvalue <- pfs_TagContainerGetString tagctr tagname
                            tagnameH <- peekCString tagname
                            tagvalueH <- peekCString tagvalue
                            loop ((tagnameH,tagvalueH):k) else return k
    tagsl <- loop []
    let w = fromIntegral $ pfs_FrameGetWidth frame
        h = fromIntegral $ pfs_FrameGetHeight frame
    (xc,yc,zc) <- alloca $ \x -> alloca $ \y -> alloca $ \z -> do
        pfs_FrameGetXYZChannels frame x y z
        xchanp <- peek x
        ychanp <- peek y
        zchanp <- peek z
        xchan <- pfs_ChannelGetRawData xchanp
        ychan <- pfs_ChannelGetRawData ychanp
        zchan <- pfs_ChannelGetRawData zchanp
        xchan' <- newForeignPtr_ xchan
        ychan' <- newForeignPtr_ ychan
        zchan' <- newForeignPtr_ zchan
        xchanA <- A.unsafeForeignPtrToStorableArray xchan' ((0,0),(w,h))
        ychanA <- A.unsafeForeignPtrToStorableArray ychan' ((0,0),(w,h))
        zchanA <- A.unsafeForeignPtrToStorableArray zchan' ((0,0),(w,h))
        return (xchanA,ychanA,zchanA)
    pfs_deleteDOMIO domio

    return $ FPFrame tagsl xc yc zc M.empty (w) (h)

-- ! Read a frame from a PFS file
readFrameFromFile :: FilePath -> IO FPFrame
readFrameFromFile path = do
    fileh <- openFile path ReadMode
    frame <- hReadFrame fileh
    hClose fileh
    return frame

-- | Write a frame to a PFS file
writeFrameToFile :: FilePath -> FPFrame -> IO ()
writeFrameToFile path frame = do    
    fileh <- openFile path WriteMode
    hWriteFrame fileh frame
    hClose fileh

-- | Write a frame to a Handle
hWriteFrame :: Handle -> FPFrame -> IO ()
hWriteFrame fileH frame = do
    domio <- pfs_newDOMIO
    fileFd <- fromIntegral <$> handleToFd fileH
    fileh <- withCString "wb" $ \filemode -> fdopen fileFd filemode

    frameh <- pfs_DOMIOCreateFrame domio (fromIntegral $ width frame) (fromIntegral $ height frame)
    A.withStorableArray (channelX frame) $ \xchan -> A.withStorableArray (channelY frame) $ \ychan -> A.withStorableArray (channelZ frame) $ \zchan ->
        alloca $ \xchanPP -> alloca $ \ychanPP -> alloca $ \zchanPP -> do
            pfs_FrameCreateXYZChannels frameh xchanPP ychanPP zchanPP
            xchanP <- peek xchanPP
            ychanP <- peek ychanPP
            zchanP <- peek zchanPP
            ptrx <- pfs_ChannelGetRawData xchanP
            copyBytes ptrx xchan ((width frame)*(height frame)*4) 
            ptry <- pfs_ChannelGetRawData ychanP
            copyBytes ptry ychan ((width frame)*(height frame)*4) 
            ptrz <- pfs_ChannelGetRawData zchanP
            copyBytes ptrz zchan ((width frame)*(height frame)*4) 
        

    forM_ (M.toList . channels $ frame) $ \(name,val) -> do 
       channel <- withCString name $ \nameC -> pfs_FrameCreateChannel frameh nameC
       ptr' <- pfs_ChannelGetRawData channel
       A.withStorableArray val $ \ptr0 -> copyBytes ptr' ptr0 ((width frame)*(height frame)*4) 

    tagctr <- pfs_FrameGetTags frameh
    forM_ (tags frame) $ \(tag,value) -> withCString tag $ \tagC -> withCString value $ \valueC -> pfs_TagContainerSetString tagctr tagC valueC

    pfs_DOMIOWriteFrame domio frameh fileh
    pfs_deleteDOMIO domio

-- | get a frame from stdin
getFrame :: IO FPFrame
getFrame = hReadFrame stdin

-- | put a frame to stdout
putFrame :: FPFrame -> IO ()
putFrame = hWriteFrame stdout