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

module PerceptualHash ( imgHash
                      , fileHash
                      , hammingDistance
                      ) where

import           Codec.Avif                    (decode)
import qualified Codec.Picture                 as JuicyPixels
import           Codec.Picture.WebP            (decodeRgb8)
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 :: forall a. Bits a => a -> a -> Int
hammingDistance a
x a
y = forall a. Bits a => a -> Int
popCount (a
x forall a. Bits a => a -> a -> a
`xor` a
y)

dct32 :: (Floating e, Array arr Y e) => Image arr Y e
dct32 :: forall e arr. (Floating e, Array arr Y e) => Image arr Y e
dct32 = forall arr cs e.
Array arr cs e =>
(Int, Int) -> ((Int, Int) -> Pixel cs e) -> Image arr cs e
makeImage (Int
32,Int
32) (Int, Int) -> Pixel Y e
gen
    where gen :: (Int, Int) -> Pixel Y e
gen (Int
i,Int
j) = forall e. e -> Pixel Y e
PixelY forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a
sqrt(e
2forall a. Fractional a => a -> a -> a
/e
n) forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
cos((forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
i forall a. Num a => a -> a -> a
* (Int
jforall a. Num a => a -> a -> a
-Int
1)) forall a. Num a => a -> a -> a
* forall a. Floating a => a
pi)forall a. Fractional a => a -> a -> a
/e
n)
          n :: e
n = e
32

idMat :: (Fractional e, Array arr X e) => Image arr X e
idMat :: forall e arr. (Fractional e, Array arr X e) => Image arr X e
idMat = forall arr cs e.
Array arr cs e =>
(Int, Int) -> ((Int, Int) -> Pixel cs e) -> Image arr cs e
makeImage (Int
7,Int
7) (\(Int, Int)
_ -> forall e. e -> Pixel X e
PixelX (e
1forall a. Fractional a => a -> a -> a
/e
49))

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

{-# INLINE size32 #-}
size32 :: Array arr cs e => Image arr cs e -> Image arr cs e
size32 :: forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
size32 = forall method arr cs e.
(Interpolation method, Array arr cs e) =>
method
-> Border (Pixel cs e)
-> (Int, Int)
-> Image arr cs e
-> Image arr cs e
resize Bilinear
Bilinear forall px. Border px
Edge (Int
32,Int
32)

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

medianImmut :: (Ord e, Fractional e, V.Vector v e) => v e -> e
medianImmut :: forall e (v :: * -> *).
(Ord e, Fractional e, Vector v e) =>
v e -> e
medianImmut v e
v = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e, Fractional e) =>
v (PrimState m) e -> m e
median forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
V.thaw v e
v

dct :: (Floating e, Array arr Y e) => Image arr Y e -> Image arr Y e
dct :: forall e arr.
(Floating e, Array arr Y e) =>
Image arr Y e -> Image arr Y e
dct Image arr Y e
img = forall e arr. (Floating e, Array arr Y e) => Image arr Y e
dct32 forall arr cs e.
Array arr cs e =>
Image arr cs e -> Image arr cs e -> Image arr cs e
|*| Image arr Y e
img forall arr cs e.
Array arr cs e =>
Image arr cs e -> Image arr cs e -> Image arr cs e
|*| forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
transpose forall e arr. (Floating e, Array arr Y e) => Image arr Y e
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 :: forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash = forall (v :: * -> *). Vector v Bool => v Bool -> Word64
asWord64 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e (v :: * -> *).
(Fractional e, Vector v e, Vector v Bool, Ord e) =>
v e -> v Bool
aboveMed forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (\(PixelY e
x) -> e
x) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e.
Array arr cs e =>
Image arr cs e -> Vector arr (Pixel cs e)
toVector forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
crop8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e arr.
(Floating e, Array arr Y e) =>
Image arr Y e -> Image arr Y e
dct forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
size32 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall e arr cs.
(Fractional e, Array arr X e, Array arr cs e) =>
Image arr cs e -> Image arr cs e
meanFilter

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

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

{-# INLINE fileWebp #-}
fileWebp :: FilePath -> IO (Image VS RGB Word8)
fileWebp :: String -> IO (Image VS RGB Word8)
fileWebp String
fp = do
    ByteString
contents <- String -> IO ByteString
BS.readFile String
fp
    let (JuicyPixels.Image Int
m Int
n Vector (PixelBaseComponent PixelRGB8)
pixels) = ByteString -> Image PixelRGB8
decodeRgb8 ByteString
contents
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall arr cs e.
Array arr cs e =>
(Int, Int) -> Vector arr (Pixel cs e) -> Image arr cs e
fromVector (Int
n, Int
m) forall a b. (a -> b) -> a -> b
$ forall a b. (Storable a, Storable b) => Vector a -> Vector b
VS.unsafeCast Vector (PixelBaseComponent PixelRGB8)
pixels

{-# INLINE readWebp #-}
readWebp :: FilePath -> IO (Image VS Y Double)
readWebp :: String -> IO (Image VS Y Double)
readWebp = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall cs e cs' e' arr.
(Convertible cs e, ToYA cs' e', ToRGBA cs' e', Array arr cs' e',
 Array arr cs e) =>
Image arr cs' e' -> Image arr cs e
convert forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS RGB Word8)
fileWebp

fileHashWebp :: FilePath -> IO Word64
fileHashWebp :: String -> IO Word64
fileHashWebp = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image VS Y Double -> Image RSU Y Double
convRepa) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS Y Double)
readWebp
    -- faster
    where convRepa :: Image VS Y Double -> Image RSU Y Double
convRepa = forall r cs e.
Source r (Pixel cs e) =>
Array r DIM2 (Pixel cs e) -> Image RSU cs e
fromRepaArrayS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e.
Array arr cs e =>
Image arr cs e -> Array U DIM2 (Pixel cs e)
toRepaArray

fileHash :: FilePath -> IO (Either String Word64)
fileHash :: String -> IO (Either String Word64)
fileHash String
fp | String
".webp" forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` String
fp = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO Word64
fileHashWebp String
fp
            | String
".avif" forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` String
fp = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO Word64
fileHashAvif String
fp
            | Bool
otherwise = String -> IO (Either String Word64)
fileHashHip String
fp

{-# INLINE readAvif #-}
readAvif :: FilePath -> IO (Image VS Y Double)
readAvif :: String -> IO (Image VS Y Double)
readAvif = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall cs e cs' e' arr.
(Convertible cs e, ToYA cs' e', ToRGBA cs' e', Array arr cs' e',
 Array arr cs e) =>
Image arr cs' e' -> Image arr cs e
convert forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS RGBA Word8)
fileAvif

fileHashAvif :: FilePath -> IO Word64
fileHashAvif :: String -> IO Word64
fileHashAvif = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image VS Y Double -> Image RSU Y Double
convRepa) forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS Y Double)
readAvif
    where convRepa :: Image VS Y Double -> Image RSU Y Double
convRepa = forall r cs e.
Source r (Pixel cs e) =>
Array r DIM2 (Pixel cs e) -> Image RSU cs e
fromRepaArrayS forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e.
Array arr cs e =>
Image arr cs e -> Array U DIM2 (Pixel cs e)
toRepaArray

{-# INLINE fileAvif #-}
fileAvif :: FilePath -> IO (Image VS RGBA Word8)
fileAvif :: String -> IO (Image VS RGBA Word8)
fileAvif String
fp = do
    (JuicyPixels.Image Int
m Int
n Vector (PixelBaseComponent PixelRGBA8)
pixels) <- ByteString -> Image PixelRGBA8
decode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO ByteString
BS.readFile String
fp
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall arr cs e.
Array arr cs e =>
(Int, Int) -> Vector arr (Pixel cs e) -> Image arr cs e
fromVector (Int
n, Int
m) forall a b. (a -> b) -> a -> b
$ forall a b. (Storable a, Storable b) => Vector a -> Vector b
VS.unsafeCast Vector (PixelBaseComponent PixelRGBA8)
pixels

fileHashHip :: FilePath -> IO (Either String Word64)
fileHashHip :: String -> IO (Either String Word64)
fileHashHip = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash :: Image RSU Y Double -> Word64)) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall arr cs e.
(Array VS cs e, Array arr cs e,
 Readable (Image VS cs e) InputFormat) =>
String -> IO (Either String (Image arr cs e))
readImage