content: {-# LANGUAGE BangPatterns
             , FlexibleContexts
             , CPP
             , GeneralizedNewtypeDeriving
             , ScopedTypeVariables
             , FlexibleInstances
             , MultiWayIf #-}
  -- | Contour tracing of binary images (zero ~ background, nonzero ~ object).
  --
  -- Terminology:
  --
  -- A binary image is an image in which the pixel is boolean (represented
  -- here  using 'Grey' and zero or non-zero pixel values).
  --
  -- All zero-value pixels are part of the "background".
  --
  -- An object is an connected group of non-zero pixels.
  --
  -- A 'Contour' is a trace of an objects' outer or inner edges.  Some
  -- objects are solid, having no inner contours (consider a filled circle,
  -- or letters such as 'h', 's', 'k' and 'l').  Other objects have "holes", also known as inner
  -- contours.  The letters 'a' and 'e' have one hole while the letter 'B' has two.
  --
  -- After obtaining a 'Contours' structure (via the 'contours' function) the
  -- raw traces ('Contour' type) can be used for further processing or the contours can be
  -- filtered by aspects of interest and selectively re-drawn ('drawContour') , perhaps used to
  -- mask the original image.
  --
  -- About Holes:
  --
  -- In cases where there is only one hole it is uniquely recorded in the
  -- 'Contours' structure.  Objects with more than one hole record all inner
  -- contours in one vector making them hard to extract separately - this is
  -- due to the main data structure not being rich enough to record the holes
  -- separately. As of writing, this is not seen as an issue because the
  -- desired operation, drawContour, can still be achieved.  Changing this
  -- behavior should be trivial if desired.
  --
  -- Use:
  --
  -- To use this library it is advised that you preprocess the image,
  -- including thresholding (ex: 'otsu' on a grey scale image), to obtain a binary image then call:
  --
  -- @
  -- cs = contours img
  -- @
  --
  -- The 'Contours' structure can be accessed directly if desired.  It
  -- includes an 'Map' of all contours (numbered counting from 1) and
  -- a vector of the contour sizes (indexed by contour number, zero index is
  -- unused/zero).
  --
  -- The algorithm implemented in this module follows the design laid out in
  -- 'A Linear-Time Component-Labeling Algorithm Using Contour Tracing Technique' [1].
  --
  -- [1] http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.95.6330&rep=rep1&type=pdf
  module Vision.Image.Contour (
      -- * Main Interface
        Contours(..), ContourId, OneContour, ContourValue, Contour(..), RowContour
      , contours
      -- * ADT style interface (hides 'Contours' internals)
      , allContourIds, lookupContour, rowContour, contourSize, contourPerimeter
      -- * Reconstructing Image Elements
      , ContourDrawStyle(..), drawContour, drawContours
      ) where
  
  #if __GLASGOW_HASKELL__ < 710
  import Control.Applicative
  #endif
  
  import Control.Monad (when)
  import Control.Monad.ST
  import qualified Data.Map.Strict as Map
  import Data.Map.Strict (Map)
  import Data.Maybe (catMaybes)
  import Data.List (groupBy,sort)
  import Data.Function (on)
  import qualified Data.Vector.Unboxed as VU
  import qualified Data.Vector.Unboxed.Mutable as VM
  import Foreign.Storable
  
  import Vision.Image.Mutable (MutableManifest, new', write)
  import qualified Vision.Image.Mutable as Mut
  import Vision.Image.Grey (Grey,GreyPixel)
  import Vision.Image.Type (Delayed)
  import Vision.Image.Class (
        MaskedImage (..), Image (..), FromFunction (..), index
      )
  import Vision.Primitive (
        Z (..), (:.) (..), Point, ix2, Size
      )
  
  --------------------------------------------------------------------------------
  --  Types and ADT-Style Interface
  
  -- | Contours of an image include:
  --    * A map from contour number to outer points and negative contour number of inner contour points.
  --    * A vector of sizes for each contour for domain [1..size contourOutlines] (the zero index is meaningless)
  data Contours =
          Contours { contourOutlines :: Map ContourId Contour
                   , contourSizes    :: !(VU.Vector Int)
                   }
  
  allContourIds :: Contours -> [ContourId]
  allContourIds = Map.keys . contourOutlines
  
  contourPerimeter :: Contours -> ContourId -> [Point]
  contourPerimeter m i =
      maybe [] (map fst . VU.toList . outerContour) (lookupContour m i)
  
  contourSize :: Contours -> ContourId -> Int
  contourSize (Contours _ s) i
      | unCID i < 0 || unCID i >= VU.length s = 0
      | otherwise                             = s VU.! unCID i
  
  lookupContour :: Contours -> ContourId -> Maybe Contour
  lookupContour m i = Map.lookup i (contourOutlines m)
  
  -- |Contours are identified by a numeric ID number.
  newtype ContourId = CID { unCID :: Int } deriving (Eq, Ord, Storable, Num, Show)
  
  -- |A contour is described by the points on the perimeter and a boolean
  -- indicating if that point is "terminal" (next pixel to
  -- the right is background iff the point is terminal).  The terminal
  -- information allows for a slightly simpler 'drawContour' implementation.
  type OneContour    = VU.Vector ContourValue
  type ContourValue  = (Point,Bool)
  data Contour = Contour { outerContour  :: OneContour
                         , innerContours :: [OneContour]
                         } -- Pair of outer and inner contours
  
  insOuterContour :: ContourId -> OneContour -> Map ContourId Contour
                                             -> Map ContourId Contour
  insOuterContour cid o mp =
      let c = Contour o []
      in Map.insert cid c mp
  
  insInnerContour :: ContourId -> OneContour -> Map ContourId Contour
                                             -> Map ContourId Contour
  insInnerContour cid i mp =
      let c = Contour (error "Impossible: Inner contour with no outer!") [i]
          f _ (Contour o is) = Contour o (i:is)
      in Map.insertWith f cid c mp
  
  -- |RowContour is a method of expressing contours by, for each row,
  -- recording the start of an object and the end (due to reaching the other
  -- side or a hole/inner contour) for each row.
  type RowContour = VU.Vector (Point,Point)
  
  --------------------------------------------------------------------------------
  --  Image Reconstruction
  
  -- | Outline: Just draw the edge.
  --
  -- OuterOutline: Outline the outer contours only, no hole contours
  -- AllOutlines: Draw all contours
  -- Fill: Draw the object but fill it in, ignoring holes.
  -- FillWithHoles: Draw the object and do not fill in the holes.
  data ContourDrawStyle = OuterOutline | AllOutlines | Fill | FillWithHoles
        deriving (Eq, Ord, Show, Read, Enum, Bounded)
  
  -- | Draws a given contour. The size specified must be large enough to
  -- include the coordinate originally occupied by the contour being drawn,
  -- no cropping or other transformation is done.
  drawContour :: Contours -> Size -> ContourDrawStyle -> ContourId -> Grey
  drawContour master sz sty c = drawContours master sz sty [c]
  
  -- |Draws many contours.  See 'drawContour'.
  drawContours :: Contours -> Size -> ContourDrawStyle -> [ContourId] -> Grey
  drawContours m sz AllOutlines ids  = drawOutlines listOfUVec m ids sz
   where listOfUVec (Contour o is) = o:is
  drawContours m sz OuterOutline ids = drawOutlines listOfUVec m ids sz
   where listOfUVec (Contour o _) = [o]
  drawContours m sz sty ids = drawRows pnts sz
   where lk = lookupContour m
         pnts = case sty of
                    Fill          -> map (VU.toList . outerContour) $ catMaybes $ map lk ids -- map (map (\(a,x) -> (a,not x)) . lk) innerIds ++ map lk outerIds
                    FillWithHoles -> map  (concatMap VU.toList . maybe [] (\x -> outerContour x : innerContours x) . lk) ids
                    _             -> error "Impossible: Style is not Fill, FillWithHoles"
  
  drawOutlines :: (Contour -> [VU.Vector ContourValue]) -> Contours -> [ContourId] -> Size -> Grey
  drawOutlines oper m ids sz = runST f
   where
    f = do
      i <- new' sz 0 :: ST s (MutableManifest GreyPixel s)
      let vs = map fst $ concatMap VU.toList $ concatMap oper $ catMaybes $ map (lookupContour m) ids
      mapM_ (\p -> write i p 255) vs
      Mut.unsafeFreeze i
  
  
  -- | Draws rows, throwing an exception when the size is too small for the
  -- coordinates.
  drawRows :: [[ContourValue]] -> Size -> Grey
  drawRows vs sz = runST $ do
      i <- new' sz 0
      mapM_ (drawMutable i) vs
      Mut.unsafeFreeze i
  
  drawMutable :: MutableManifest GreyPixel s -> [ContourValue] -> ST s ()
  drawMutable i cs = VU.mapM_ (f i) rs
   where
       rs = rowContour cs
       f img (start,stop) = go (start, stop)
         where go (s@(Z:.row:.col),t) = do
                  write img s 255
                  when (s /= t) $ go (Z :. row :. (col+1),t)
  
  -- |Given a vector including outer (and optionally inner) contour points,
  -- make 'row contour' from which is easier to transform back into a binary
  -- image.  By not including the inner contour points the row will be filled, making
  -- traces of objects with holes appear solid.
  rowContour :: [ContourValue] -> RowContour
  rowContour cs =
      let rows :: [[(Point,Bool)]]
          rows = groupBy ((==) `on` ((\(Z:.r:._) -> r) . fst)) $ sort cs -- XXX consider vector quick/tim sort
      in VU.fromList $ concatMap walkM rows
   where
    walkM :: [(Point,Bool)] -> [(Point,Point)]
    walkM [x] = [(fst x,fst x)]
    walkM x   = maybe (error $ "Impossible: No terminal when walking contour: " ++ show (x,cs)) id $ walk x
    walk :: [(Point,Bool)] -> Maybe [(Point,Point)]
    walk [] = Just []
    walk xs@(x:_) = case dropWhile (not . snd) xs of
                        []     -> Nothing
                        (t:ys) -> ((fst x,fst t) :) <$> walk ys
  
  -- |The meat of this module is the 'contours' function, which extracts
  -- the contours (outer and inner outlines) of a binary image.
  -- Zero-valued pixels are the background and non-zero are active/objects to
  -- trace.  The output, 'Contours', contains enough information to determine
  -- the number of contours, their traces, the size in pixels (filled size
  -- and perimeter), number of holes, etc.
  contours :: (Image src, Num (ImagePixel src), Eq (ImagePixel src)) => src -> Contours
  contours src = runST $ do
       let bsrc = fromFunction (Z :. y+2 :. x+2) mkBorder
       mutImg   <- new' (shape bsrc) zid
       (outlines,sz) <- doLabeling bsrc mutImg
       sizes <- freezeBlobSizes sz
       return (Contours outlines sizes)
   where
   (Z :. y :. x) = shape src
  
   mkBorder (Z :. j :. i)
     | j == 0 || j == (y+1) || i == 0 || i == (x+1) = background
     | otherwise                                    = index src (Z :. j-1 :. i-1)
  
  -- The image is assumed to be binary and should have values of either 0 (black) or... nonzero (white)
  -- here we assume the background is black.  Nonzero would require more
  -- change elsewhere!
  background :: Num a => a
  background = 0
  
  zid :: ContourId
  zid = CID 0
  
  data BlobSizes s = BS (VM.MVector s Int)
  
  freezeBlobSizes :: BlobSizes s -> ST s (VU.Vector Int)
  freezeBlobSizes (BS v) = VU.unsafeFreeze v
  
  incBlobSizes :: ContourId -> BlobSizes s -> ST s (BlobSizes s)
  incBlobSizes (CID i) s@(BS v)
    | i > 0 =
       if VM.length v <= i
           then do nv <- VM.unsafeGrow v (i*2)
                   mapM_ (\ix -> VM.unsafeWrite nv ix 0) [i..i*2-1]
                   VM.unsafeWrite nv i 1
                   return (BS nv)
           else do p <- VM.unsafeRead v i
                   VM.unsafeWrite v i (p+1)
                   return s
    | otherwise = return s
  
  zeroBlobSizes :: ST s (BlobSizes s)
  zeroBlobSizes = BS <$> VM.replicate 1024 0
  
  -- Make a contour image of the same dimension but with ContourIDs instead
  -- of pixels.
  doLabeling :: forall s p. (Storable p, Num p, Eq p) => Delayed p -> MutableManifest ContourId s -> ST s (Map ContourId Contour,BlobSizes s)
  doLabeling src mutImg = zeroBlobSizes >>= go (Just $ ix2 1 1) (CID 0) (CID 1) mempty
   where
   getCID    :: Point -> ST s ContourId
   getCID     = Mut.read mutImg
   setCID i c = write mutImg i c
   getPixel :: Point -> ImagePixel (Delayed p)
   getPixel   = index src
  
   incIx :: Point -> Maybe Point
   incIx !(Z :. (!y) :. (!x))
      | x < xMax-1  = Just $ Z :. y     :. (x+1)
      | y < yMax-1  = Just $ Z :. (y+1) :. 1
      | otherwise = Nothing
  
   (Z :. yMax :. xMax) = shape src
  
   -- Traverse the source image top to bottom, left to right.  If the pixel
   -- has an ID then propagate that ID to all the following active pixels in
   -- the row.  If the pixel is active and has no ID then trace either an
   -- inner or outer contour.  If the pixel is inactive then skip it.
   go Nothing   _ _ !mp v              = return (mp,v)
   go (Just idx) leftCID !newCID !mp v =
     do thisCID <- getCID idx
        if | val == background                     -> skipForward -- this step doesn't appear in the paper! D'oh
           | thisCID == zid && above == background ->
                           do -- Step 1: Outer contour trace (active pixel with id=0 and above is background)
                              newContour <- traceContour src mutImg ExternalContour idx newCID
                              go (Just idx) newCID (newCID + 1) (insOuterContour newCID newContour mp) v
           | below == background ->               -- Step 2: P is an              white pixel
               do belowCID <- getCID belowIdx     --                 ^          ^
                  if | belowCID == zid ->         --                 ^ unmarked ^
                           do -- Step 2a: Inner contour trace, below pixel was unmarked
                              let innerCID = if zid == thisCID then leftCID else thisCID
                              inner <- traceContour src mutImg InternalContour idx innerCID
                              go (incIx idx) innerCID newCID (insInnerContour innerCID inner mp) v
                              -- there can be more than one inner contour, make a richer container structure than IntMap?
                              -- Notice this isn't entirely necessary, one
                              -- CID can contain all inner contours and they can be redrawn correctly.
                     | otherwise -> stepForward -- Step 2b: Previously-observed contour
           | otherwise                             -> stepForward -- Active pixel not on a contour
     where val         = getPixel idx
           above       = getPixel (Z :. y-1 :. x)
           below       = getPixel belowIdx
           belowIdx    = Z :. y+1 :. x
           Z :. y :. x = idx
           stepForward = do xId <- if leftCID <= zid
                                    then getCID idx
                                    else return leftCID
                            setCID idx xId
                            nv <- incBlobSizes xId v
                            go (incIx idx) xId newCID mp nv
           skipForward = go (incIx idx) (-2) newCID mp v
  
  -- Mark surrounding background pixels
  -- label non-background pixels with CID
  --
  -- Unroll the loop one step to account for the lone-pixel case.  Without
  -- lone pixels the tight inner loop can save a check (See 'Impossible')
  --
  -- TODO: optimize later, duplicate tracer and remove the p==pos, after a benchmarking method is setup.
  traceContour :: forall p s. (Storable p, Num p, Eq p) => Delayed p -> MutableManifest ContourId s -> ContourType -> Point -> ContourId -> ST s OneContour
  traceContour src mutImg contourTy origPnt assignedCID =
    do next <- tracer origPnt startPos
       case next of
           Nothing              -> return (VU.fromList $ fixList [(origPnt,True)])
           Just (sndPnt,sndPos) -> do
              let f pnt pos = do (nPnt,nPos) <- maybe (error "Impossible: Nothing in inner") id <$> tracer pnt pos
                                 if pnt == origPnt && nPnt == sndPnt
                                     then return [] -- XXX some algorithms duplicate the start point, pnt, as the last point `return [pnt]`.  Should we?
                                     else ((pnt,terminal pnt) :) <$> f nPnt nPos
              VU.fromList . fixList . ((origPnt, terminal origPnt):) <$> f sndPnt sndPos
  
   where
     terminal (Z :. row :. col) = 0 == getPixel (Z :. row :. (col+1))
     -- Translate between indexes in our border-added image and the original
     fixList xs = let f (Z :. a :. b, t) = (Z :. a-1 :. b-1,t) in map f xs
     startPos   = case contourTy of { ExternalContour -> UR ; InternalContour -> LL  }
     setCID i c = write mutImg i c
     getPixel :: Point -> ImagePixel (Delayed p)
     getPixel   = index src
  
     {-# INLINE tracer #-}
     tracer pnt pos =
         let tracer' True p | p == pos = setCID pnt assignedCID >> return Nothing
             tracer' _ p = do let rpnt = relPoint pnt p
                                  v    = getPixel rpnt
                              if | v == background -> do setCID rpnt (-1)
                                                         tracer' True (incCP p)
                                 | otherwise       -> do setCID pnt assignedCID
                                                         return (Just (rpnt, decCP2 p))
         in tracer' False pos
  
  
  --------------------------------------------------------------------------------
  --  Internal Types and Utilities
  
  data ContourType = ExternalContour | InternalContour deriving (Eq)
  
  -- A contour position is Upper/Lower/Middle Left/Center/Right pixel
  -- relative to the current.  Because we add a border to the image prior
  -- to processing, all original pixels `cross` contours positions are
  -- a valid point.
  data ContourPos  = MR | LR | LC | LL
                   | ML | UL | UC | UR
              deriving (Enum, Bounded, Eq, Show)
  
  relPoint :: Point -> ContourPos -> Point
  relPoint (Z :. row :. col) pos = Z :. row' :. col'
   where !row' = row + y
         !col' = col + x
         x = colOffset VU.! fromEnum pos
         y = rowOffset VU.! fromEnum pos
  
  -- Rather than have branching (and expected poor performance),
  -- use a table of x/y offsets for the relative pixel position.
  colOffset,rowOffset :: VU.Vector Int
  rowOffset = VU.fromList [0,1,1,1,0,-1,-1,-1]
  colOffset = VU.fromList [1,1,0,-1,-1,-1,0,1]
  
  -- Position clockwise by one tick
  incCP :: ContourPos -> ContourPos
  incCP  = toEnum . ((`rem` 8) . (+ 1)) . fromEnum
  
  -- Position counter-clockwise by two ticks
  decCP2 :: ContourPos -> ContourPos
  decCP2 = toEnum . ((`rem` 8) . (+ 6)) . fromEnum
  