{-# OPTIONS_GHC -fspecialize-aggressively #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies     #-}

module PerceptualHash ( imgHash
                      , fileHash
                      , hammingDistance
                      ) where

#ifdef AVIF
import           Codec.Avif                    (decode)
#endif
import qualified Codec.Picture                 as JuicyPixels
#ifdef WEBP
import           Codec.Picture.WebP            (decodeRgb8)
#endif
import           Control.Monad.ST              (runST)
import           Data.Bits                     (Bits, popCount, shiftL, xor, (.|.))
import qualified Data.ByteString               as BS
import           Data.List                     (isSuffixOf)
import qualified Data.Vector.Generic           as V
import qualified Data.Vector.Storable          as VS
import           Data.Word                     (Word64, Word8)
import           Graphics.Image                (Array, Bilinear (..), Border (Edge, Reflect), Image,
                                                Pixel (PixelX, PixelY), RGB, RGBA, RSU (..), VS, X,
                                                Y, convert, convolve, crop, makeImage, readImage,
                                                resize, transpose, (|*|))
import           Graphics.Image.Interface      (fromVector, toVector)
import qualified Graphics.Image.Interface      as Hip
import           Graphics.Image.Interface.Repa (fromRepaArrayS, toRepaArray)
import           Median                        (median)

-- | See
-- [wiki](https://en.wikipedia.org/wiki/Hamming_distance#Algorithm_example).
--
-- @since 0.1.4.0
{-# SPECIALIZE hammingDistance :: Word64 -> Word64 -> Int #-}
hammingDistance :: Bits a => a -> a -> Int
hammingDistance x y = popCount (x `xor` y)

dct32 :: (Floating e, Array arr Y e) => Image arr Y e
dct32 = makeImage (32,32) gen
    where gen (i,j) = PixelY $ sqrt(2/n) * cos((fromIntegral (i * (j-1)) * pi)/n)
          n = 32

idMat :: (Fractional e, Array arr X e) => Image arr X e
idMat = makeImage (7,7) (\_ -> PixelX (1/49))

{-# INLINE meanFilter #-}
meanFilter :: (Fractional e, Array arr X e, Array arr cs e) => Image arr cs e -> Image arr cs e
meanFilter = {-# SCC "meanFilter" #-} convolve Reflect idMat

{-# INLINE size32 #-}
size32 :: Array arr cs e => Image arr cs e -> Image arr cs e
size32 = resize Bilinear Edge (32,32)

crop8 :: Array arr cs e => Image arr cs e -> Image arr cs e
crop8 = crop (0,0) (8,8)

medianImmut :: (Ord e, Fractional e, V.Vector v e) => v e -> e
medianImmut v = runST $
    median =<< V.thaw v

dct :: (Floating e, Array arr Y e) => Image arr Y e -> Image arr Y e
dct img = dct32 |*| img |*| transpose dct32

{-# INLINE imgHash #-}
-- | DCT based hash. See
-- [Zauner](https://www.phash.org/docs/pubs/thesis_zauner.pdf).
--
-- It is suggested that you use this with the Repa backend.
imgHash :: (Ord e, Floating e, Array arr Y e, Array arr X e, V.Vector (Hip.Vector arr) Bool, V.Vector (Hip.Vector arr) e) => Image arr Y e -> Word64
imgHash = asWord64 . aboveMed . V.map (\(PixelY x) -> x) . toVector . crop8 . dct . size32 . meanFilter

asWord64 :: V.Vector v Bool => v Bool -> Word64
asWord64 = V.foldl' (\acc x -> (acc `shiftL` 1) .|. boolToWord64 x) 0
    where boolToWord64 :: Bool -> Word64
          boolToWord64 False = 0
          boolToWord64 True  = 1

aboveMed :: (Fractional e, V.Vector v e, V.Vector v Bool, Ord e) => v e -> v Bool
aboveMed v =
    let med = medianImmut v
    in V.map (<med) v

#ifdef WEBP
{-# INLINE fileWebp #-}
fileWebp :: FilePath -> IO (Image VS RGB Word8)
fileWebp fp = do
    contents <- BS.readFile fp
    let (JuicyPixels.Image m n pixels) = decodeRgb8 contents
    pure $ fromVector (n, m) $ VS.unsafeCast pixels

{-# INLINE readWebp #-}
readWebp :: FilePath -> IO (Image VS Y Double)
readWebp = fmap convert . fileWebp

fileHashWebp :: FilePath -> IO Word64
fileHashWebp = fmap (imgHash . convRepa) . readWebp
    -- faster
    where convRepa = fromRepaArrayS . toRepaArray
#endif

fileHash :: FilePath -> IO (Either String Word64)
fileHash fp
#ifdef WEBP
    | ".webp" `isSuffixOf` fp = pure <$> fileHashWebp fp
#endif
#ifdef AVIF
    | ".avif" `isSuffixOf` fp = pure <$> fileHashAvif fp
#endif
    | otherwise = fileHashHip fp

#ifdef AVIF
{-# INLINE readAvif #-}
readAvif :: FilePath -> IO (Image VS Y Double)
readAvif = fmap convert . fileAvif

fileHashAvif :: FilePath -> IO Word64
fileHashAvif = fmap (imgHash . convRepa) . readAvif
    where convRepa = fromRepaArrayS . toRepaArray

{-# INLINE fileAvif #-}
fileAvif :: FilePath -> IO (Image VS RGBA Word8)
fileAvif fp = do
    (JuicyPixels.Image m n pixels) <- decode <$> BS.readFile fp
    pure $ fromVector (n, m) $ VS.unsafeCast pixels
#endif

fileHashHip :: FilePath -> IO (Either String Word64)
fileHashHip = fmap (fmap (imgHash :: Image RSU Y Double -> Word64)) . readImage