{-#LANGUAGE TypeSynonymInstances#-}
 -- | Image thresholding operations
module CV.Thresholding(
  -- * Interfaces to OpenCV functions
  ThresholdType(..)
, threshold
, thresholdOtsu
, AdaptiveType(..)
, adaptiveThreshold
  -- * Other methods
, bernsen
, nibbly
, nibblyr
, kittler
, kittlerMeasure
, betweenClassVariance
) where

import CV.Image
import CV.Filters
import qualified CV.ImageMath as IM
import CV.ImageMathOp
import CV.Morphology
import System.IO.Unsafe
import CV.Sampling
import Utils.List
import Data.List
import CV.Histogram
import CV.Bindings.ImgProc
import Foreign.Ptr (nullPtr,castPtr)

-- | A type class for selecting the maximum value for each image type, used in
--   creating the thresholded image
class MaxVal a where
  maxval :: (Image GrayScale a) -> Double

instance MaxVal D32 where
  maxval _ = 1

instance MaxVal D8 where
  maxval _ = 255

-- | Thresholding behavior for values larger and smaller than threshold
data ThresholdType
  -- | Values larger than threshold are set to max, smaller to zero
  = MaxAndZero
  -- | Values larger than threshold are set to zero, smaller to max
  | ZeroAndMax
  -- | Values larger than threshold are truncated to threshold, smaller are not touched
  | ThreshAndValue
  -- | Values larger than threshold are not touched, smaller are set to zero
  | ValueAndZero
  -- | Values larger than threshold are set to zero, smaller are not touched
  | ZeroAndValue

-- | Utility function for converting ThresholdType to c values
cThresholdType t =
  case t of
    MaxAndZero     -> c'CV_THRESH_BINARY
    ZeroAndMax     -> c'CV_THRESH_BINARY_INV
    ThreshAndValue -> c'CV_THRESH_TRUNC
    ValueAndZero   -> c'CV_THRESH_TOZERO
    ZeroAndValue   -> c'CV_THRESH_TOZERO_INV

-- | Utility function for converting otsu ThresholdType to c values
cOtsuThresholdType t =
  case t of
    MaxAndZero     -> c'CV_THRESH_OTSU_BINARY
    ZeroAndMax     -> c'CV_THRESH_OTSU_BINARY_INV
    ThreshAndValue -> c'CV_THRESH_OTSU_TRUNC
    ValueAndZero   -> c'CV_THRESH_OTSU_TOZERO
    ZeroAndValue   -> c'CV_THRESH_OTSU_TOZERO_INV

-- | Method used for selecting the adaptive threshold value
data AdaptiveType
  -- | Threshold using the arithmetic mean of pixel neighborhood
  = ByMean
  -- | Threshold using the gaussian weighted mean of pixel neighborhood
  | ByGaussian 

-- | Utility function for converting AdaptiveType to c value
cAdaptiveType t =
  case t of
    ByMean     -> c'CV_ADAPTIVE_THRESH_MEAN_C
    ByGaussian -> c'CV_ADAPTIVE_THRESH_GAUSSIAN_C

-- | Thresholds a grayscale image according to the selected type, using the
--   given threshold value.
threshold :: (MaxVal d) => ThresholdType -> Double -> Image GrayScale d -> Image GrayScale d
threshold ttype tval image =
  unsafePerformIO $
    withCloneValue image $ \result ->
      withImage image $ \pimage ->
        withImage result $ \presult -> do
          c'cvThreshold (castPtr pimage) (castPtr presult) (realToFrac tval)
            (realToFrac (maxval image)) (cThresholdType ttype)
          return result

-- | Thresholds a grayscale image using the otsu method according to the
--   selected type. Threshold value is selected automatically, and only 8-bit
--   images are supported.
thresholdOtsu :: ThresholdType -> Image GrayScale D8 -> Image GrayScale D8
thresholdOtsu ttype image =
  unsafePerformIO $
    withCloneValue image $ \result ->
      withImage image $ \pimage ->
        withImage result $ \presult -> do
          c'cvThreshold (castPtr pimage) (castPtr presult) 0
            (realToFrac 255) (cOtsuThresholdType ttype)
          return result

-- | Applies adaptive thresholding by selecting the optimal threshold value for
--   each pixel. The threshold is selected by calculating the arithmetic or
--   gaussian-weighted mean of a pixel neighborhood, and applying a bias term to
--   the obtained value.
adaptiveThreshold :: (MaxVal d) => AdaptiveType -> ThresholdType -> Int -> Double 
  -> Image GrayScale d -> Image GrayScale d
adaptiveThreshold a t neighborhood bias image =
  unsafePerformIO $
    withCloneValue image $ \result ->
      withImage image $ \pimage ->
        withImage result $ \presult -> do
          c'cvAdaptiveThreshold (castPtr pimage) (castPtr presult) (realToFrac (maxval image))
            (cAdaptiveType a) (cThresholdType t) (fromIntegral neighborhood) (realToFrac bias)
          return result

bernsen (w,h) c i = goodContrast #* (i #< surface)
            where
             low  = erode se 1  i
             high = dilate se 1 i
             goodContrast = IM.moreThan c (high #- low)
             surface = 0.5 |* (high #+ low) 
             se = structuringElement (w,h) (w`div`2,h`div`2) EllipseShape

-- Very slow implementation of niblack thresholding
--niblack (w,h) k i = IM.more2Than trunc (unsafePerformIO $ surface) 
--    where
--     trunc = getRegion (w`div`2,h`div`2) (wi-w,hi-h) i
--     (wi,hi) = getSize i
--     surface = renderFlatList (wi-w,hi-h) (map th patches)
--     th ptch = IM.average ptch + k * IM.stdDeviation ptch
--     patches = allPatches (w,h) i

nibbly k c i = let dev = IM.stdDeviation i
                   mean = IM.average i 
             in IM.moreThan (mean+k*dev+c) i

nibblyr (w,h) k i = IM.lessThan t flat
    where
     t = IM.average flat + k * IM.stdDeviation flat
     flat = i #- gaussian (w,h) i


---- TODO: Convert Histograms from Doubles to Floats..
--otsu bs image = IM.moreThan (realToFrac threshold) image
--    where
--        histogram  = getHistogram bs $ image
--        partitions = histogramPartitions histogram 
--        (threshold,_,_) = maximumBy (comparing otsuCmp) partitions 
--        otsuCmp (t,as,bs) = betweenClassVariance (as) (bs)

-- This is excruciatingly slow means of finding kittler-illingworth threshold
-- for an image
kittler precision image = IM.moreThan t image
    where t = maximumBy (comparing (kittlerMeasure image))
                            [0,0+precision..1]

kittlerMeasure image t = unNaN $ 
                        p_t*log fgDev
                      + (1-p_t)*log bgDev
                      - p_t*log p_t 
                      - (1-p_t)*log(1-p_t)
    where
     unNaN x | isNaN x = -10000000
             | otherwise = x
     thresholded = unsafeImageTo32F (IM.lessThan t image)
     p_t = IM.sum ( thresholded) / fromIntegral (getArea image)
     bgDev = realToFrac $ IM.stdDeviationMask image thresholded 
     fgDev = realToFrac $ IM.stdDeviationMask image (IM.invert thresholded)


--histogramPartitions (HGD a) = zip3 (head.tails.map fst $ a) 
--                                   (tail.inits.map snd $ a) 
--                                   (reverse.tail.reverse.tails.map snd $ a)

betweenClassVariance as bs = sum as * sum bs 
                             * (average bs - average as)^2