{-# LANGUAGE TypeOperators #-}
module MXNet.Coco.Mask where

import Data.Word
import qualified Data.ByteString as BS
import Data.Array.Repa (Array, Z(..), (:.)(..), DIM1, DIM2, DIM3, extent)
import Data.Array.Repa.Repr.Unboxed
import qualified Data.Vector.Unboxed as UV (convert, length)
import qualified Data.Vector.Storable as SV (Vector, map, unsafeCast)

import MXNet.Coco.Raw

-- mask should be of 3 dimension and in CWH order
type Mask = Array U DIM3 Word8
type Area = Array U DIM1 Word32
type Iou  = Array U DIM2 Double
type BBox = Array U DIM2 Double
type Poly = SV.Vector Double

data CompactRLE = CompactRLE Int Int [BS.ByteString]

encode :: Mask -> IO CompactRLE
encode mask = do
    let Z :. n :. w :. h = extent mask
    -- assuming Word8 are identical with CUChar
    rles <- rleEncode (SV.map fromIntegral $ UV.convert $ toUnboxed mask) h w n
    CompactRLE h w <$> mapM rleToString rles

decode :: CompactRLE -> IO Mask
decode im@(CompactRLE h w bss) = do
    let n = length bss
    rles <- frString im
    raw <- rleDecode rles h w
    return $ 
        fromUnboxed (Z :. n :. w :. h) $
        UV.convert $
        SV.map fromIntegral raw

merge :: CompactRLE -> Bool -> IO CompactRLE
merge im intersect = do
    let CompactRLE h w bss = im
        n = length bss
    if n > 1 then do
        rles <- frString im
        orle <- rleMerge rles intersect
        bs <- rleToString orle
        return $ CompactRLE h w [bs]
    else 
        return im

area :: CompactRLE -> IO Area
area im@(CompactRLE _ _ bss) = do
    let num = length bss
    rles <- frString im
    as <- rleArea rles num
    -- assuming the area can be represented as Word32
    return $ 
        fromListUnboxed (Z :. num) $
        map fromIntegral $ as

iouRLEs :: CompactRLE -> CompactRLE -> [Bool] -> IO Iou
iouRLEs dt gt iscrowd = do
    dt_rles <- frString dt
    gt_rles <- frString gt
    ((m, n), arr) <- rleIou dt_rles gt_rles iscrowd
    return $ fromListUnboxed (Z :. n :. m) arr

iouBBs :: BBox -> BBox -> [Bool] -> IO Iou
iouBBs bb1 bb2 iscrowd = do
    let bb1' = BB $ SV.unsafeCast $ UV.convert $ toUnboxed bb1
        bb2' = BB $ SV.unsafeCast $ UV.convert $ toUnboxed bb2
    ((m, n), arr) <- bbIou bb1' bb2' iscrowd
    return $ fromListUnboxed (Z :. n :. m) arr

toBBox :: CompactRLE -> IO BBox
toBBox im = do
    rles <- frString im
    BB bb <- rleToBbox rles
    let bb' = UV.convert $ SV.unsafeCast bb
    return $ fromUnboxed (Z :. UV.length bb' :. 4) bb'

frBBox :: BBox -> Int -> Int -> IO CompactRLE
frBBox bb h w = do
    rles <- rleFrBbox (BB $ SV.unsafeCast $ UV.convert $ toUnboxed bb) h w
    CompactRLE h w <$> mapM rleToString rles

frPoly :: [Poly] -> Int -> Int -> IO CompactRLE
frPoly polys h w = do
    rles <- mapM (\poly -> rleFrPoly (SV.unsafeCast poly) h w) polys
    CompactRLE h w <$> mapM rleToString rles

frUncompressedRLE :: [Int] -> Int -> Int -> IO CompactRLE
frUncompressedRLE raw h w = do
    orle <- rleInit h w (map fromIntegral raw)
    crle <- rleToString orle
    return $ CompactRLE h w [crle]

frString :: CompactRLE -> IO [RLE]
frString (CompactRLE h w bss) = mapM (\bs -> rleFrString bs h w) bss