{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
-- | This module provide some color quantisation algorithm
-- in order to help in the creation of paletted images.
-- The most important function is `palettize` which will
-- make everything to create a nice color indexed image
-- with its palette.
module Codec.Picture.ColorQuant
    ( palettize
    , defaultPaletteOptions
    , PaletteCreationMethod(..)
    , PaletteOptions( .. )
    ) where

#if !MIN_VERSION_base(4,8,0)
import           Control.Applicative (Applicative (..), (<$>))
#endif

import           Data.Bits           (unsafeShiftL, unsafeShiftR, (.&.), (.|.))
import           Data.List           (elemIndex)
import           Data.Maybe          (fromMaybe)
import           Data.Set            (Set)
import qualified Data.Set            as Set
import           Data.Word           (Word32)

import           Data.Vector         (Vector, (!))
import qualified Data.Vector         as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Storable as VS

import           Codec.Picture.Types

-------------------------------------------------------------------------------
----            Palette Creation and Dithering
-------------------------------------------------------------------------------

-- | Define which palette creation method is used.
data PaletteCreationMethod =
      -- | MedianMeanCut method, provide the best results (visualy)
      -- at the cost of increased calculations.
      MedianMeanCut
      -- | Very fast algorithm (one pass), doesn't provide good
      -- looking results.
    | Uniform

-- | To specify how the palette will be created.
data PaletteOptions = PaletteOptions
    { -- | Algorithm used to find the palette
      paletteCreationMethod :: PaletteCreationMethod

      -- | Do we want to apply the dithering to the
      -- image. Enabling it often reduce compression
      -- ratio but enhance the perceived quality
      -- of the final image.
    , enableImageDithering  :: Bool

      -- | Maximum number of color we want in the
      -- palette
    , paletteColorCount     :: Int
    }

-- | Default palette option, which aim at the best quality
-- and maximum possible colors (256)
defaultPaletteOptions :: PaletteOptions
defaultPaletteOptions = PaletteOptions
    { paletteCreationMethod = MedianMeanCut
    , enableImageDithering  = True
    , paletteColorCount     = 256
    }

-- | Reduces an image to a color palette according to `PaletteOpts` and
--   returns the /indices image/ along with its `Palette`.
palettize :: PaletteOptions -> Image PixelRGB8 -> (Image Pixel8, Palette)
palettize opts@PaletteOptions { paletteCreationMethod = method } =
  case method of
    MedianMeanCut -> medianMeanCutQuantization opts
    Uniform       -> uniformQuantization opts

-- | Modified median cut algorithm with optional ordered dithering. Returns an
-- image of `Pixel8` that acts as a matrix of indices into the `Palette`.
medianMeanCutQuantization :: PaletteOptions -> Image PixelRGB8
                          -> (Image Pixel8, Palette)
medianMeanCutQuantization opts img
  | isBelow =
      (pixelMap okPaletteIndex img, vecToPalette okPaletteVec)
  | enableImageDithering opts = (pixelMap paletteIndex dImg, palette)
  | otherwise = (pixelMap paletteIndex img, palette)
  where
    maxColorCount = paletteColorCount opts
    (okPalette, isBelow) = isColorCountBelow maxColorCount img
    okPaletteVec = V.fromList $ Set.toList okPalette
    okPaletteIndex p = nearestColorIdx p okPaletteVec

    palette = vecToPalette paletteVec
    paletteIndex p = nearestColorIdx p paletteVec
    paletteVec = mkPaletteVec cs
    cs =  Set.toList . clusters maxColorCount $ img
    dImg = pixelMapXY dither img

-- | A naive one pass Color Quantiation algorithm - Uniform Quantization.
-- Simply take the most significant bits. The maxCols parameter is rounded
-- down to the nearest power of 2, and the bits are divided among the three
-- color channels with priority order green, red, blue. Returns an
-- image of `Pixel8` that acts as a matrix of indices into the `Palette`.
uniformQuantization :: PaletteOptions -> Image PixelRGB8 -> (Image Pixel8, Palette)
uniformQuantization opts img
  -- -| colorCount img <= maxCols = colorQuantExact img
  | enableImageDithering opts =
        (pixelMap paletteIndex (pixelMapXY dither img), palette)
  | otherwise = (pixelMap paletteIndex img, palette)
  where
    maxCols = paletteColorCount opts
    palette = listToPalette paletteList
    paletteList = [PixelRGB8 r g b | r <- [0,dr..255]
                                   , g <- [0,dg..255]
                                   , b <- [0,db..255]]
    (bg, br, bb) = bitDiv3 maxCols
    (dr, dg, db) = (2^(8-br), 2^(8-bg), 2^(8-bb))
    paletteIndex (PixelRGB8 r g b) = fromIntegral $ fromMaybe 0 (elemIndex
      (PixelRGB8 (r .&. (255 - dr)) (g .&. (255 - dg)) (b .&. (255 - db)))
      paletteList)

isColorCountBelow :: Int -> Image PixelRGB8 -> (Set.Set PixelRGB8, Bool)
isColorCountBelow maxColorCount img = go 0 Set.empty
  where rawData = imageData img
        maxIndex = VS.length rawData
        
        go !idx !allColors
            | Set.size allColors > maxColorCount = (Set.empty, False)
            | idx >= maxIndex - 2 = (allColors, True)
            | otherwise = go (idx + 3) $ Set.insert px allColors
                where px = unsafePixelAt rawData idx 

vecToPalette :: Vector PixelRGB8 -> Palette
vecToPalette ps = generateImage (\x _ -> ps ! x) (V.length ps) 1

listToPalette :: [PixelRGB8] -> Palette
listToPalette ps = generateImage (\x _ -> ps !! x) (length ps) 1

bitDiv3 :: Int -> (Int, Int, Int)
bitDiv3 n = case r of
            0 -> (q, q, q)
            1 -> (q+1, q, q)
            _ -> (q+1, q+1, q)
  where
    r = m `mod` 3
    q = m `div` 3
    m = floor . logBase (2 :: Double) $ fromIntegral n

-------------------------------------------------------------------------------
----            Dithering
-------------------------------------------------------------------------------

-- Add a dither mask to an image for ordered dithering.
-- Uses a small, spatially stable dithering algorithm based on magic numbers
-- and arithmetic inspired by the /a dither/ algorithm of Øyvind Kolås,
-- pippin@gimp.org, 2013. See, http://pippin.gimp.org/a_dither/.
dither :: Int -> Int -> PixelRGB8 -> PixelRGB8
dither x y (PixelRGB8 r g b) = PixelRGB8 (fromIntegral r')
                                         (fromIntegral g')
                                         (fromIntegral b')
  where
    -- Should view 16 as a parameter that can be optimized for best looking
    -- results
    r' = min 255 (fromIntegral r + (x' + y') .&. 16)
    g' = min 255 (fromIntegral g + (x' + y' + 7973) .&. 16)
    b' = min 255 (fromIntegral b + (x' + y' + 15946) .&. 16)
    x' = 119 * x
    y' = 28084 * y

-------------------------------------------------------------------------------
----            Small modification of foldl package by Gabriel Gonzalez
-------------------------------------------------------------------------------

-- Modification to Control.foldl by Gabriel Gonzalez copyright 2013, BSD3.
-- http://hackage.haskell.org/package/foldl-1.0.1/docs/Control-Foldl.html

{-| Efficient representation of a left fold that preserves the fold's step
    function, initial accumulator, and extraction function

    This allows the 'Applicative' instance to assemble derived folds that
    traverse the container only once
-}
data Fold a b = forall x . Fold (x -> a -> x) x (x -> b)

{-| Apply a strict left 'Fold' to a 'Foldable' container

    Much slower than 'fold' on lists because 'Foldable' operations currently do
    not trigger @build/foldr@ fusion
-}
fold :: Fold PackedRGB b -> VU.Vector PackedRGB -> b
fold (Fold step begin done) = done . VU.foldl' step begin
{-# INLINE fold #-}

{-
F.foldr :: (a -> b -> b) -> b -> t a -> b

fold :: (Foldable f) => Fold a b -> f a -> b
fold (Fold step begin done) as = F.foldr step' done as begin
  where step' x k z = k $! step z x
-}

data Pair a b = Pair !a !b

instance Functor (Fold a) where
    fmap f (Fold step begin done) = Fold step begin (f . done)
    {-# INLINABLE fmap #-}

instance Applicative (Fold a) where
    pure b    = Fold (\() _ -> ()) () (\() -> b)
    {-# INLINABLE pure #-}
    (Fold stepL beginL doneL) <*> (Fold stepR beginR doneR) =
        let step (Pair xL xR) a = Pair (stepL xL a) (stepR xR a)
            begin = Pair beginL beginR
            done (Pair xL xR) = doneL xL $ doneR xR
        in  Fold step begin done
    {-# INLINABLE (<*>) #-}

{- | Like 'length', except with a more general 'Num' return value -}
intLength :: Fold a Int
intLength = Fold (\n _ -> n + 1) 0 id

-------------------------------------------------------------------------------
----            Modified Median Cut Algorithm
-------------------------------------------------------------------------------

-- Based on the OCaml implementation:
-- http://rosettacode.org/wiki/Color_quantization
-- which is in turn based on: www.leptonica.com/papers/mediancut.pdf.
-- We use the product of volume and population to determine the next cluster
-- to split and determine the placement of each color by compating it to the
-- mean of the parent cluster. So median cut is a bit of a misnomer, since one
-- of the modifiations is to use the mean.

mkPaletteVec :: [Cluster] -> Vector PixelRGB8
mkPaletteVec  = V.fromList . map (toRGB8 . meanColor)

type PackedRGB = Word32

data Cluster = Cluster
    { value       :: {-# UNPACK #-} !Float
    , meanColor   :: !PixelRGBF
    , dims        :: !PixelRGBF
    , colors      :: VU.Vector PackedRGB
    }

instance Eq Cluster where
    a == b =
        (value a, meanColor a, dims a) == (value b, meanColor b, dims b)

instance Ord Cluster where
    compare a b =
        compare (value a, meanColor a, dims a) (value b, meanColor b, dims b)

data Axis = RAxis | GAxis | BAxis

inf :: Float
inf = read "Infinity"

fromRGB8 :: PixelRGB8 -> PixelRGBF
fromRGB8 (PixelRGB8 r g b) =
  PixelRGBF (fromIntegral r) (fromIntegral g) (fromIntegral b)

toRGB8 :: PixelRGBF -> PixelRGB8
toRGB8 (PixelRGBF r g b) =
  PixelRGB8 (round r) (round g) (round b)

meanRGB :: Fold PixelRGBF PixelRGBF
meanRGB = mean <$> intLength <*> pixelSum
  where
    pixelSum = Fold (mixWith $ const (+)) (PixelRGBF 0 0 0) id
    mean n = colorMap (/ nf)
      where nf = fromIntegral n

minimal :: Fold PixelRGBF PixelRGBF
minimal = Fold mini (PixelRGBF inf inf inf) id
  where mini = mixWith $ const min

maximal :: Fold PixelRGBF PixelRGBF
maximal = Fold maxi (PixelRGBF (-inf) (-inf) (-inf)) id
  where maxi = mixWith $ const max

extrems :: Fold PixelRGBF (PixelRGBF, PixelRGBF)
extrems = (,) <$> minimal <*> maximal

volAndDims :: Fold PixelRGBF (Float, PixelRGBF)
volAndDims = deltify <$> extrems
  where deltify (mini, maxi) = (dr * dg * db, delta)
          where delta@(PixelRGBF dr dg db) =
                        mixWith (const (-)) maxi mini

unpackFold :: Fold PixelRGBF a -> Fold PackedRGB a
unpackFold (Fold step start done) = Fold (\acc -> step acc . transform) start done
  where transform = fromRGB8 . rgbIntUnpack

mkCluster :: VU.Vector PackedRGB -> Cluster
mkCluster ps = Cluster
    { value = v * fromIntegral l
    , meanColor = m
    , dims = ds
    , colors = ps
    }
  where
    worker = (,,) <$> volAndDims <*> meanRGB <*> intLength
    ((v, ds), m, l) = fold (unpackFold worker) ps

maxAxis :: PixelRGBF -> Axis
maxAxis (PixelRGBF r g b) =
  case (r `compare` g, r `compare` b, g `compare` b) of
    (GT, GT, _)  -> RAxis
    (LT, GT, _)  -> GAxis
    (GT, LT, _)  -> BAxis
    (LT, LT, GT) -> GAxis
    (EQ, GT, _)  -> RAxis
    (_,  _,  _)  -> BAxis

-- Split a cluster about its largest axis using the mean to divide up the
-- pixels.
subdivide :: Cluster -> (Cluster, Cluster)
subdivide cluster = (mkCluster px1, mkCluster px2)
  where
    (PixelRGBF mr mg mb) = meanColor cluster
    (px1, px2) = VU.partition (cond . rgbIntUnpack) $ colors cluster
    cond = case maxAxis $ dims cluster of
      RAxis -> \(PixelRGB8 r _ _) -> fromIntegral r < mr
      GAxis -> \(PixelRGB8 _ g _) -> fromIntegral g < mg
      BAxis -> \(PixelRGB8 _ _ b) -> fromIntegral b < mb

rgbIntPack :: PixelRGB8 -> PackedRGB
rgbIntPack (PixelRGB8 r g b) =
    wr `unsafeShiftL` (2 * 8) .|. wg `unsafeShiftL` 8 .|. wb
  where wr = fromIntegral r
        wg = fromIntegral g
        wb = fromIntegral b

rgbIntUnpack :: PackedRGB -> PixelRGB8
rgbIntUnpack v = PixelRGB8 r g b
  where
    r = fromIntegral $ v `unsafeShiftR` (2 * 8)
    g = fromIntegral $ v `unsafeShiftR` 8
    b = fromIntegral v

initCluster :: Image PixelRGB8 -> Cluster
initCluster img = mkCluster $ VU.generate ((w * h) `div` subSampling) packer
  where samplingFactor = 3
        subSampling = samplingFactor * samplingFactor
        compCount = componentCount (undefined :: PixelRGB8)
        w = imageWidth img
        h = imageHeight img
        rawData = imageData img
        packer ix =
            rgbIntPack . unsafePixelAt rawData $ ix * subSampling * compCount

-- Take the cluster with the largest value = (volume * population) and remove it
-- from the priority queue. Then subdivide it about its largest axis and put the
-- two new clusters on the queue.
split :: Set Cluster -> Set Cluster
split cs = Set.insert c1 . Set.insert c2  $ cs'
  where
    (c, cs') = Set.deleteFindMax cs
    (c1, c2) = subdivide c

-- Keep splitting the initial cluster until there are 256 clusters, then return
-- a priority queue containing all 256.
clusters :: Int -> Image PixelRGB8 -> Set Cluster
clusters maxCols img = clusters' (maxCols - 1)
  where
    clusters' :: Int -> Set Cluster
    clusters' 0 = Set.singleton c
    clusters' n = split (clusters' (n-1))
    c = initCluster img

-- Euclidean distance squared, between two pixels.
dist2Px :: PixelRGB8 -> PixelRGB8 -> Int
dist2Px (PixelRGB8 r1 g1 b1) (PixelRGB8 r2 g2 b2) = dr*dr + dg*dg + db*db
  where
    (dr, dg, db) =
      ( fromIntegral r1 - fromIntegral r2
      , fromIntegral g1 - fromIntegral g2
      , fromIntegral b1 - fromIntegral b2 )

nearestColorIdx :: PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx p ps  = fromIntegral $ V.minIndex (V.map (`dist2Px` p) ps)