{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Graphics.Image.IO.Histogram
-- Copyright   : (c) Alexey Kuleshevich 2017
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.Image.IO.Histogram (
  Histogram(..), Histograms, ChannelColour(..), getHistograms, getHistogram,
  equalizeHistogram, cdf,
  displayHistograms, writeHistograms
  ) where

import           Control.Concurrent                        (forkIO)
import           Control.Monad                             (void)
import qualified Data.Colour                               as C
import qualified Data.Colour.Names                         as C
import qualified Data.Vector.Unboxed                       as V
import           Prelude                                   as P
import           System.Directory                          (getTemporaryDirectory)
import           System.FilePath                           ((</>))
import           System.IO.Temp                            (createTempDirectory)

import           Graphics.Image.ColorSpace
import           Graphics.Image.Interface                  as I
import           Graphics.Image.IO
import           Graphics.Rendering.Chart.Backend.Diagrams
import           Graphics.Rendering.Chart.Easy

#if MIN_VERSION_vector(0,11,0)
import           Data.Vector.Unboxed.Mutable               (modify)
#else
import           Control.Monad.Primitive                   (PrimMonad (..))
import qualified Data.Vector.Unboxed.Mutable               as MV

modify :: (PrimMonad m, V.Unbox a) => MV.MVector (PrimState m) a -> (a -> a) -> Int -> m ()
modify v f idx = do
  e <- MV.read v idx
  MV.write v idx $ f e
#endif


class ChannelColour cs where

  -- | Get a pure colour representation of a channel.
  csColour :: cs -> C.AlphaColour Double


-- | A single channel histogram of an image.
data Histogram = Histogram { hBins   :: V.Vector Int
                             -- ^ Vector containing pixel counts. Index of a
                             -- vector serves as an original pixel value.
                           , hName   :: String
                             -- ^ Name of the channel that will be displayed in
                             -- the legend.
                           , hColour :: C.AlphaColour Double
                             -- ^ Color of a plotted line.
                           }
-- | For now it is just a type synonym, but in the future it might become a custom
-- data type with fields like title, width, heigth, etc.
type Histograms = [Histogram]


-- | Create a histogram per channel with 256 bins each.
getHistograms :: forall arr cs e . (ChannelColour cs, MArray arr X e, Array arr X e,
                                    MArray arr cs e, Array arr cs e) =>
                 Image arr cs e
              -> Histograms
getHistograms = P.zipWith setCh (enumFrom (toEnum 0) :: [cs]) . P.map getHistogram . toImagesX
  where setCh !cs !h = h { hName = show cs
                         , hColour = csColour cs }
        {-# INLINE setCh #-}
{-# INLINE getHistograms #-}

-- | Generate a histogram with 256 bins for a single channel Gray image.
getHistogram :: MArray arr X e =>
                Image arr X e
             -> Histogram
getHistogram !img = Histogram { hBins = V.modify countBins $
                                        V.replicate
                                        (1 + fromIntegral (maxBound :: Word8)) (0 :: Int)
                              , hName = show X
                              , hColour = csColour X }
  where
    incBin !v (PixelX x) = modify v (+1) $ fromIntegral (toWord8 x)
    {-# INLINE incBin #-}
    countBins !v = I.mapM_ (incBin v) img
    {-# INLINE countBins #-}
{-# INLINE getHistogram #-}


-- | Discrete cumulative distribution function.
cdf :: MArray arr X e => Image arr X e -> V.Vector Double
cdf !img = V.unfoldr gen (0, 0)
  where
    !p = 1 / fromIntegral (m * n)
    !(m, n) = dims img
    !histogram = hBins $ getHistogram img
    !l = V.length histogram
    gen !(acc, k)
      | k == l = Nothing
      | otherwise =
        let !acc' = acc + fromIntegral (V.unsafeIndex histogram k)
        in Just (acc' * p, (acc', k + 1))
    {-# INLINE gen #-}
{-# INLINE cdf #-}


-- | Converts an image to Luma and performs histogram equalization.
equalizeHistogram
  :: ( ToY cs e
     , Array arr cs e
     , Array arr Y Double
     , Array arr X Word8
     , MArray arr X Word8
     )
  => Image arr cs e -> Image arr Y Double
equalizeHistogram !img = I.map f imgX where
  !imgX = I.map (toX . toPixelY) img
  !cdfImg = cdf imgX
  toX (PixelY y) = PixelX $ toWord8 y
  {-# INLINE toX #-}
  f (PixelX x) = PixelY (V.unsafeIndex cdfImg (fromIntegral x))
  {-# INLINE f #-}
{-# INLINE equalizeHistogram #-}


-- | Write histograms into a PNG image file.
--
-- >>> frog <- readImageRGB VU "images/frog.jpg"
-- >>> writeHistograms "images/frog_histogram.svg" $ getHistograms frog
--
-- <<images/frog_histogram.svg>>
--
writeHistograms :: FilePath -> Histograms -> IO ()
writeHistograms fileName hists =
  toFile def fileName $ do
    layout_title .= "Histogram"
    setColors $ P.map hColour hists
    let axis = set la_nTicks 20 . set la_nLabels 14
    layout_x_axis . laxis_generate .= scaledIntAxis (axis defaultIntAxis) (0, 260)
    let plotHist h = plot $ line (hName h) [V.toList . V.imap (,) $ hBins h]
    P.mapM_ plotHist hists

-- | Display image histograms using an external program. Works in a similar way as
-- `Graphics.Image.IO.displayImage`.
--
-- >>> frog <- readImageRGB VU "images/frog.jpg"
-- >>> displayHistograms $ getHistograms frog
--
displayHistograms :: Histograms -> IO ()
displayHistograms = displayHistogramsUsing defaultViewer False


-- | Display image histograms using an external program. Works in a similar way as
-- `Graphics.Image.IO.displayImageUsing`.
displayHistogramsUsing :: ExternalViewer
                       -> Bool
                       -> Histograms -> IO ()
displayHistogramsUsing viewer block hists = do
  let display = do
        tmpDir <- getTemporaryDirectory
        histPath <- fmap (</> "tmp-hist.svg") (createTempDirectory tmpDir "hip-histogram")
        writeHistograms histPath hists
        displayImageFile viewer histPath
  if block
    then display
    else void $ forkIO display

instance ChannelColour X where
  csColour _ = C.opaque C.darkgray

instance ChannelColour Y where
  csColour _ = C.opaque C.darkgray

instance ChannelColour YA where
  csColour LumaYA  = csColour LumaY
  csColour AlphaYA = C.opaque C.gray


instance ChannelColour RGB where
  csColour RedRGB   = C.opaque C.red
  csColour GreenRGB = C.opaque C.green
  csColour BlueRGB  = C.opaque C.blue

instance ChannelColour RGBA where
  csColour RedRGBA   = C.opaque C.red
  csColour GreenRGBA = C.opaque C.green
  csColour BlueRGBA  = C.opaque C.blue
  csColour AlphaRGBA = C.opaque C.gray


instance ChannelColour HSI where
  csColour HueHSI = C.opaque C.purple
  csColour SatHSI = C.opaque C.orange
  csColour IntHSI = C.opaque C.darkblue

instance ChannelColour HSIA where
  csColour HueHSIA   = C.opaque C.purple
  csColour SatHSIA   = C.opaque C.orange
  csColour IntHSIA   = C.opaque C.darkblue
  csColour AlphaHSIA = C.opaque C.gray


instance ChannelColour CMYK where
  csColour CyanCMYK = C.opaque C.cyan
  csColour MagCMYK  = C.opaque C.magenta
  csColour YelCMYK  = C.opaque C.yellow
  csColour KeyCMYK  = C.opaque C.black

instance ChannelColour CMYKA where
  csColour CyanCMYKA  = csColour CyanCMYK
  csColour MagCMYKA   = csColour MagCMYK
  csColour YelCMYKA   = csColour YelCMYK
  csColour KeyCMYKA   = csColour KeyCMYK
  csColour AlphaCMYKA = C.opaque C.grey


instance ChannelColour YCbCr where
  csColour LumaYCbCr  = C.opaque C.darkgray
  csColour CBlueYCbCr = C.opaque C.darkblue
  csColour CRedYCbCr  = C.opaque C.darkred


instance ChannelColour YCbCrA where
  csColour LumaYCbCrA  = csColour LumaYCbCr
  csColour CBlueYCbCrA = csColour CBlueYCbCr
  csColour CRedYCbCrA  = csColour CRedYCbCr
  csColour AlphaYCbCrA = C.opaque C.gray