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

import Prelude as P 
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 System.Directory (getTemporaryDirectory)
import System.FilePath ((</>))
import System.IO.Temp (createTempDirectory)

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

#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 Gray e, Array arr Gray 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 . toGrayImages
  where setCh cs h = h { hName = show cs
                       , hColour = csColour cs }

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

-- | 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 Gray 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