{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      :  Codec.Picture.Extra
-- Copyright   :  © 2016–present Mark Karpov
-- License     :  BSD 3 clause
--
-- Maintainer  :  Mark Karpov <markkarpov92@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- A collection of functions to scale, crop, flip images with JuicyPixels.
module Codec.Picture.Extra
  ( -- * Scaling
    scaleBilinear,

    -- * Cropping
    crop,
    trim,

    -- * Rotation
    flipHorizontally,
    flipVertically,
    rotateLeft90,
    rotateRight90,
    rotate180,

    -- * Other
    beside,
    below,
    square,
  )
where

import Codec.Picture
import qualified Codec.Picture.Types as M
import Control.Monad.ST
import Data.List (find, foldl1')
import Data.Maybe (fromMaybe)

-- | Scale an image using bi-linear interpolation.
scaleBilinear ::
  ( Pixel a,
    Bounded (PixelBaseComponent a),
    Integral (PixelBaseComponent a)
  ) =>
  -- | Desired width
  Int ->
  -- | Desired height
  Int ->
  -- | Original image
  Image a ->
  -- | Scaled image
  Image a
scaleBilinear width height img@Image {..}
  | width <= 0 || height <= 0 =
      generateImage (error "scaleBilinear: absurd") (max 0 width) (max 0 height)
  | otherwise = runST $ do
      mimg <- M.newMutableImage width height
      let sx, sy :: Float
          sx = fromIntegral imageWidth / fromIntegral width
          sy = fromIntegral imageHeight / fromIntegral height
          go x' y'
            | x' >= width = go 0 (y' + 1)
            | y' >= height = M.unsafeFreezeImage mimg
            | otherwise = do
                let xf = fromIntegral x' * sx
                    yf = fromIntegral y' * sy
                    x, y :: Int
                    x = floor xf
                    y = floor yf
                    δx = xf - fromIntegral x
                    δy = yf - fromIntegral y
                    pixelAt' i j =
                      pixelAt img (min (pred imageWidth) i) (min (pred imageHeight) j)
                writePixel mimg x' y' $
                  mulp (pixelAt' x y) ((1 - δx) * (1 - δy))
                    `addp` mulp (pixelAt' (x + 1) y) (δx * (1 - δy))
                    `addp` mulp (pixelAt' x (y + 1)) ((1 - δx) * δy)
                    `addp` mulp (pixelAt' (x + 1) (y + 1)) (δx * δy)
                go (x' + 1) y'
      go 0 0
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelRGBA16 -> Image M.PixelRGBA16 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelRGBA8 -> Image M.PixelRGBA8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelCMYK16 -> Image M.PixelCMYK16 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelCMYK8 -> Image M.PixelCMYK8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelYCbCr8 -> Image M.PixelYCbCr8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelRGB16 -> Image M.PixelRGB16 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelYCbCrK8 -> Image M.PixelYCbCrK8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelRGB8 -> Image M.PixelRGB8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelYA16 -> Image M.PixelYA16 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.PixelYA8 -> Image M.PixelYA8 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.Pixel32 -> Image M.Pixel32 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.Pixel16 -> Image M.Pixel16 #-}
{-# SPECIALIZE scaleBilinear :: Int -> Int -> Image M.Pixel8 -> Image M.Pixel8 #-}

mulp :: (Pixel a, Integral (PixelBaseComponent a)) => a -> Float -> a
mulp pixel x = colorMap (floor . (* x) . fromIntegral) pixel
{-# INLINE mulp #-}

addp ::
  forall a.
  ( Pixel a,
    Bounded (PixelBaseComponent a),
    Integral (PixelBaseComponent a)
  ) =>
  a ->
  a ->
  a
addp = mixWith (const f)
  where
    f x y =
      fromIntegral $
        (maxBound :: PixelBaseComponent a) `min` (fromIntegral x + fromIntegral y)
{-# INLINE addp #-}

-- | Crop an image. If the supplied coordinates are greater than the size of
-- the image, the image boundaries are used instead.
crop ::
  Pixel a =>
  -- | Index (X axis) of first pixel to include
  Int ->
  -- | Index (Y axis) of first pixel to include
  Int ->
  -- | Width of resulting image
  Int ->
  -- | Height of resulting image
  Int ->
  -- | Original image
  Image a ->
  -- | Cropped image
  Image a
crop x' y' w' h' img@Image {..} =
  generateImage gen w h
  where
    gen i j = pixelAt img (x + i) (y + j)
    x = min (imageWidth - 1) x'
    y = min (imageHeight - 1) y'
    w = min (imageWidth - x) w'
    h = min (imageHeight - y) h'
{-# INLINEABLE crop #-}

-- | Trim the completely transparent edges of an image.
--
-- @since 0.6.0
trim :: (Pixel a, Eq (PixelBaseComponent a)) => Image a -> Image a
trim img@Image {..} = crop left top width height img
  where
    isInvisible p = pixelOpacity p == 0
    isInvisibleRow y = all isInvisible $ flip (pixelAt img) y <$> [0 .. imageWidth - 1]
    isInvisibleCol x = all isInvisible $ pixelAt img x <$> [0 .. imageHeight - 1]

    top = fromMaybe 0 (find (not . isInvisibleRow) [0 .. imageHeight - 1])
    bottom = fromMaybe 0 (find (not . isInvisibleRow) [imageHeight - 1, imageHeight - 2 .. 0]) + 1
    height = bottom - top

    left = fromMaybe 0 (find (not . isInvisibleCol) [0 .. imageWidth - 1])
    right = fromMaybe 0 (find (not . isInvisibleCol) [imageWidth - 1, imageWidth - 2 .. 1]) + 1
    width = right - left
{-# INLINEABLE trim #-}

-- | Flip an image horizontally.
flipHorizontally :: Pixel a => Image a -> Image a
flipHorizontally img@Image {..} =
  generateImage gen imageWidth imageHeight
  where
    gen x = pixelAt img (imageWidth - 1 - x)
{-# INLINEABLE flipHorizontally #-}

-- | Flip an image vertically.
flipVertically :: Pixel a => Image a -> Image a
flipVertically img@Image {..} =
  generateImage gen imageWidth imageHeight
  where
    gen x y = pixelAt img x (imageHeight - 1 - y)
{-# INLINEABLE flipVertically #-}

-- | Rotate an image to the left by 90°.
rotateLeft90 :: Pixel a => Image a -> Image a
rotateLeft90 img@Image {..} =
  generateImage gen imageHeight imageWidth
  where
    gen x y = pixelAt img (imageWidth - 1 - y) x
{-# INLINEABLE rotateLeft90 #-}

-- | Rotate an image to the right by 90°.
rotateRight90 :: Pixel a => Image a -> Image a
rotateRight90 img@Image {..} =
  generateImage gen imageHeight imageWidth
  where
    gen x y = pixelAt img y (imageHeight - 1 - x)
{-# INLINEABLE rotateRight90 #-}

-- | Rotate an image by 180°, i.e flip both vertically and horizontally.
--
-- @since 0.2.0
rotate180 :: Pixel a => Image a -> Image a
rotate180 img@(Image w h _) = generateImage g w h
  where
    g x y = pixelAt img (w - 1 - x) (h - 1 - y)
{-# INLINEABLE rotate180 #-}

-- | Create an image by placing several images side by side. If the images
-- are of differnet heights the smallest height is used.
--
-- @since 0.2.0
beside :: Pixel a => [Image a] -> Image a
beside = foldl1' go
  where
    go :: Pixel a => Image a -> Image a -> Image a
    go img1@(Image w1 h1 _) img2@(Image w2 h2 _) =
      generateImage g (w1 + w2) h
      where
        g x
          | x < w1 = pixelAt img1 x
          | otherwise = pixelAt img2 (x - w1)
        h = min h1 h2
{-# INLINEABLE beside #-}

-- | Create an image by placing several images in a vertical stack. If the
-- images are of differnet widths the smallest width is used.
--
-- @since 0.2.0
below :: Pixel a => [Image a] -> Image a
below = foldl1' go
  where
    go :: Pixel a => Image a -> Image a -> Image a
    go img1@(Image w1 h1 _) img2@(Image w2 h2 _) =
      generateImage g w (h1 + h2)
      where
        g x y
          | y < h1 = pixelAt img1 x y
          | otherwise = pixelAt img2 x (y - h1)
        w = min w1 w2
{-# INLINEABLE below #-}

-- | Make an image a perfect square by adding filler around it.
--
-- @since 0.6.0
square :: Pixel a => a -> Image a -> Image a
square filler img@Image {..} =
  if imageWidth == imageHeight
    then img
    else generateImage gen size size
  where
    size = max imageWidth imageHeight
    extraWidth = size - imageWidth
    extraHeight = size - imageHeight
    offsetX = extraWidth `div` 2
    offsetY = extraHeight `div` 2
    gen i _ | i < offsetX = filler
    gen i _ | i >= imageWidth + offsetX = filler
    gen _ j | j < offsetY = filler
    gen _ j | j >= imageHeight + offsetY = filler
    gen i j = pixelAt img (i - offsetX) (j - offsetY)
{-# INLINEABLE square #-}