{-# LANGUAGE PackageImports, BangPatterns, QuasiQuotes, PatternGuards #-}
{-# OPTIONS -Wall -fno-warn-missing-signatures -fno-warn-incomplete-patterns #-}

-- | Canny edge detector.
--
--   NOTE: for best performance this needs to be compiled with the following GHC options:
--         -fllvm -optlo-O3 -Odph -fno-liberate-case
--         -funfolding-use-threshold100 -funfolding-keeness-factor100
--
import Data.List
import Data.Word
import Data.Int
import Control.Monad
import System.Environment
import Data.Array.Repa 				as R
import Data.Array.Repa.Stencil
import Data.Array.Repa.Specialised.Dim2
import Data.Array.Repa.IO.BMP
import Data.Array.Repa.IO.Timing
import Prelude					hiding (compare)

import System.IO.Unsafe
import qualified Data.Vector.Unboxed.Mutable	as VM
import qualified Data.Vector.Unboxed		as V
import Prelude					as P

type Image a	= Array DIM2 a

-- Constants ------------------------------------------------------------------
-- TODO: It would be better to use Word8 to represent the edge orientations,
--       but doing so currently triggers a bug in the LLVM mangler.
orientUndef	= 0	:: Float
orientPosDiag	= 64	:: Float
orientVert	= 128	:: Float
orientNegDiag	= 192	:: Float
orientHoriz	= 255	:: Float

data Edge	= None | Weak | Strong
edge None	= 0 	:: Word8
edge Weak	= 128 	:: Word8
edge Strong	= 255	:: Word8

-- Main routine ---------------------------------------------------------------
main 
 = do	args	<- getArgs
	case args of
	 [fileIn, fileOut]		
	   -> run 0 50 100 fileIn fileOut

	 [loops, threshLow, threshHigh, fileIn, fileOut]
	   -> run (read loops) (read threshLow) (read threshHigh) fileIn fileOut

	 _ -> putStrLn "repa-canny [<loops::Int> <threshLow::Int> <threshHigh::Int>] <fileIn.bmp> <fileOut.bmp>"


run loops threshLow threshHigh fileIn fileOut
 = do	arrInput 	<- liftM (force . either (error . show) id) 
			$ readImageFromBMP fileIn

	arrInput `deepSeqArray` return ()
	(arrResult, tTotal)
	 <- time
	 $ do	{-let stage str arr 
			= timeStage loops str $ return arr-}
	
		arrGrey		<- timeStage loops "toGreyScale"  $ return $ toGreyScale    arrInput
		arrBluredX	<- timeStage loops "blurX" 	  $ return $ blurSepX       arrGrey
		arrBlured	<- timeStage loops "blurY" 	  $ return $ blurSepY       arrBluredX
		arrDX		<- timeStage loops "diffX"	  $ return $ gradientX      arrBlured
		arrDY		<- timeStage loops "diffY"	  $ return $ gradientY      arrBlured
		
		arrMagOrient	<- timeStage loops "magOrient"	 
				$ return $ gradientMagOrient threshLow arrDX arrDY

		arrSuppress	<- timeStage loops "suppress"     
				$ return $ suppress threshLow threshHigh arrMagOrient

		arrStrong	<- timeStage loops "select"	$ return $ selectStrong arrSuppress	
		arrEdges	<- timeStage loops "wildfire"	$ return $ wildfire arrSuppress arrStrong	

		return arrEdges

	when (loops >= 1)
	 $ putStrLn $ "\nTOTAL\n"
	
	putStr $ prettyTime tTotal
	
	writeComponentsToBMP fileOut arrResult arrResult arrResult


-- | Wrapper to time each stage of the algorithm.
timeStage
	:: (Shape sh, Elt a)
	=> Int
	-> String 
	-> (IO (Array sh a))
	-> (IO (Array sh a))

{-# NOINLINE timeStage #-}
timeStage loops name fn
 = do	let burn !n
	     = do arr	<- fn
		  arr `deepSeqArray` return ()
		  if n <= 1 then return arr
		            else burn (n - 1)
			
	(arrResult, t)
	 <- time $ do	arrResult' <- burn loops
		   	arrResult' `deepSeqArray` return arrResult'

	when (loops >= 1) 
	 $ putStr 	$  name P.++ "\n"
			P.++ unlines [ "  " P.++ l | l <- lines $ prettyTime t ]

	return arrResult


-------------------------------------------------------------------------------
-- | The default conversions supplied via the prelude go via a GMP function
--   call instead of just using the appropriate primop.
{-# INLINE floatOfWord8 #-}
floatOfWord8 :: Word8 -> Float
floatOfWord8 w8
 	= fromIntegral (fromIntegral w8 :: Int)


{-# INLINE word8OfFloat #-}
word8OfFloat :: Float -> Word8
word8OfFloat f
 	= fromIntegral (truncate f :: Int)


-------------------------------------------------------------------------------
-- | RGB to greyscale conversion.
{-# NOINLINE toGreyScale #-}
toGreyScale :: Array DIM3 Word8 -> Array DIM2 Word8
toGreyScale 
	arr@(Array _ [Region RangeAll (GenManifest _)])
  = arr `deepSeqArray` force2
  $ unsafeTraverse arr
	(\(sh :. _) -> sh)
	(\get ix    -> rgbToLuminance 
				(get (ix :. 0))
				(get (ix :. 1))
				(get (ix :. 2)))

 where	{-# INLINE rgbToLuminance #-}
	rgbToLuminance :: Word8 -> Word8 -> Word8 -> Word8
	rgbToLuminance r g b 
	 = word8OfFloat 
		( floatOfWord8 r * 0.3
		+ floatOfWord8 g * 0.59
		+ floatOfWord8 b * 0.11)


-- | Separable Gaussian blur in the X direction.
{-# NOINLINE blurSepX #-}
blurSepX :: Array DIM2 Word8 -> Array DIM2 Float
blurSepX arr@(Array _ [Region RangeAll (GenManifest _)])	
	= arr `deepSeqArray` force2
	$ forStencilFrom2  BoundClamp arr floatOfWord8
	  [stencil2|	1 4 6 4 1 |]	


-- | Separable Gaussian blur in the Y direction.
{-# NOINLINE blurSepY #-}
blurSepY :: Array DIM2 Float -> Array DIM2 Float
blurSepY arr@(Array _ [Region RangeAll (GenManifest _)])	
	= arr `deepSeqArray` force2
	$ R.map (/ 256)
	$ forStencil2  BoundClamp arr
	  [stencil2|	1
	 		4
			6
			4
			1 |]


-- | Compute gradient in the X direction.
{-# NOINLINE gradientX #-}
gradientX :: Array DIM2 Float -> Array DIM2 Float
gradientX img@(Array _ [Region RangeAll (GenManifest _)])
 	= img `deepSeqArray` force2
    	$ forStencil2 (BoundConst 0) img
	  [stencil2|	-1  0  1
			-2  0  2
			-1  0  1 |]


-- | Compute gradient in the Y direction.
{-# NOINLINE gradientY #-}
gradientY :: Array DIM2 Float -> Array DIM2 Float
gradientY img@(Array _ [Region RangeAll (GenManifest _)])
	= img `deepSeqArray` force2
	$ forStencil2 (BoundConst 0) img
	  [stencil2|	 1  2  1
			 0  0  0
			-1 -2 -1 |] 


-- | Classify the magnitude and orientation of the vector gradient.
{-# NOINLINE gradientMagOrient #-}
gradientMagOrient :: Float -> Image Float -> Image Float -> Image (Float, Float)
gradientMagOrient !threshLow 
 	dX@(Array _ [Region RangeAll (GenManifest _)])
	dY@(Array _ [Region RangeAll (GenManifest _)])

 = [dX, dY] `deepSeqArrays` force2
 $ R.zipWith magOrient dX dY

 where	{-# INLINE magOrient #-}
	magOrient :: Float -> Float -> (Float, Float)
	magOrient !x !y
		= (magnitude x y, orientation x y)
	
	{-# INLINE magnitude #-}
	magnitude :: Float -> Float -> Float
	magnitude !x !y
		= sqrt (x * x + y * y)
	
	{-# INLINE orientation #-}
	orientation :: Float -> Float -> Float
	orientation !x !y

	 -- Don't bother computing orientation if vector is below threshold.
 	 | x >= negate threshLow, x < threshLow
 	 , y >= negate threshLow, y < threshLow
 	 = orientUndef

	 | otherwise
	 = let	-- Determine the angle of the vector and rotate it around a bit
		-- to make the segments easier to classify.
		!d	= atan2 y x 
		!dRot	= (d - (pi/8)) * (4/pi)
	
		-- Normalise angle to beween 0..8
		!dNorm	= if dRot < 0 then dRot + 8 else dRot

		-- Doing explicit tests seems to be faster than using the FP floor function.
	   in	if dNorm >= 4
		 then if dNorm >= 6
			then if dNorm >= 7
				then orientHoriz   -- 7
				else orientNegDiag -- 6

			else if dNorm >= 5
				then orientVert    -- 5
				else orientPosDiag -- 4

		 else if dNorm >= 2
			then if dNorm >= 3
				then orientHoriz   -- 3
				else orientNegDiag -- 2

			else if dNorm >= 1
				then orientVert    -- 1
				else orientPosDiag -- 0


-- | Suppress pixels that are not local maxima, and use the magnitude to classify maxima
--   into strong and weak (potential) edges.
{-# NOINLINE suppress #-}
suppress :: Float -> Float -> Image (Float, Float) -> Image Word8
suppress threshLow threshHigh
	   dMagOrient@(Array shSrc [Region RangeAll (GenManifest _)]) 
 = dMagOrient `deepSeqArray` force2 
 $ makeBordered2 shSrc 1 
		(GenCursor id addDim (const 0))
 		(GenCursor id addDim comparePts)

 where	{-# INLINE comparePts #-}
	comparePts d@(sh :. i :. j)
	 | o == orientUndef     = edge None
         | o == orientHoriz	= isMax (getMag (sh :. i   :. j-1)) (getMag (sh :. i   :. j+1)) 
         | o == orientVert	= isMax (getMag (sh :. i-1 :. j))   (getMag (sh :. i+1 :. j)) 
         | o == orientNegDiag	= isMax (getMag (sh :. i-1 :. j-1)) (getMag (sh :. i+1 :. j+1)) 
         | o == orientPosDiag	= isMax (getMag (sh :. i-1 :. j+1)) (getMag (sh :. i+1 :. j-1)) 
         | otherwise 		= edge None
      
         where
          !o 		= getOrient d  
          !m		= getMag    (Z :. i :. j)

	  getMag 	= fst . (R.unsafeIndex dMagOrient)
	  getOrient	= snd . (R.unsafeIndex dMagOrient)

	  {-# INLINE isMax #-}
          isMax intensity1 intensity2
            | m < threshLow 	= edge None
            | m < intensity1 	= edge None
            | m < intensity2 	= edge None
	    | m < threshHigh	= edge Weak
	    | otherwise		= edge Strong


-- | Select indices of strong edges.
--   TODO: If would better if we could medge this into the above stage, and record the strong edge
--         during non-maximum suppression, but Repa doesn't provide a fused mapFilter primitive yet.
{-# NOINLINE selectStrong #-}
selectStrong :: Image Word8 -> Array DIM1 Int
selectStrong img@(Array _ [Region RangeAll (GenManifest vec)])
 = img `deepSeqArray` 
   let 	{-# INLINE match #-}
	match ix	= vec `V.unsafeIndex` ix == edge Strong

	{-# INLINE process #-}
	process ix	= ix
	
   in	select match process (size $ extent img)


-- | Trace out strong edges in the final image. 
--   Also trace out weak edges that are connected to strong edges.
{-# NOINLINE wildfire #-}
wildfire 
	:: Image Word8		-- ^ Image with strong and weak edges set.
	-> Array DIM1 Int	-- ^ Array containing flat indices of strong edges.
	-> Image Word8

wildfire img@(Array _ [Region RangeAll (GenManifest _)]) arrStrong
 = unsafePerformIO 
 $ do	(sh, vec)	<- wildfireIO 
	return	$ sh `seq` vec `seq` 
		  Array sh [Region RangeAll (GenManifest vec)]

 where	lenImg		= R.size $ R.extent img
	lenStrong	= R.size $ R.extent arrStrong
	vStrong		= toVector arrStrong
	
	wildfireIO
  	 = do	-- Stack of image indices we still need to consider.
		vStrong' <- V.thaw vStrong
		vStack	 <- VM.grow vStrong' (lenImg - lenStrong)
	
		-- Burn in new edges.
		vImg	<- VM.unsafeNew lenImg
		VM.set vImg 0
		burn vImg vStack lenStrong
		vImg'	<- V.unsafeFreeze vImg
		return	(extent img, vImg')

	
	burn :: VM.IOVector Word8 -> VM.IOVector Int -> Int -> IO ()
	burn !vImg !vStack !top
	 | top == 0
	 = return ()
	
	 | otherwise
	 = do	let !top'		=  top - 1
		n			<- VM.unsafeRead vStack top'
		let (Z :. y :. x)	= fromIndex (R.extent img) n

		let {-# INLINE push #-}
		    push t		= pushWeak vImg vStack t
				
		VM.write vImg n (edge Strong)
	    	 >>  push (Z :. y - 1 :. x - 1) top'
	    	 >>= push (Z :. y - 1 :. x    )
	    	 >>= push (Z :. y - 1 :. x + 1)

	    	 >>= push (Z :. y     :. x - 1)
	    	 >>= push (Z :. y     :. x + 1)

	    	 >>= push (Z :. y + 1 :. x - 1)
	    	 >>= push (Z :. y + 1 :. x    )
	    	 >>= push (Z :. y + 1 :. x + 1)

	    	 >>= burn vImg vStack

	-- If this ix is weak in the source then set it to strong in the
	-- result and push the ix onto the stack.
	{-# INLINE pushWeak #-}
	pushWeak vImg vStack ix top
	 = do	let n		= toIndex (extent img) ix
		xDst		<- VM.unsafeRead vImg n
		let xSrc	= img `R.unsafeIndex` ix

		if   xDst == edge None 
		  && xSrc == edge Weak
		 then do
			VM.unsafeWrite vStack top (toIndex (extent img) ix)
			return (top + 1)
			
		 else	return top