{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
module Graphics.Rasterific.Shading
    ( transformTextureToFiller
    , sampledImageShader
    , plotOpaquePixel
    , unsafePlotOpaquePixel
    ) where

import Control.Monad.ST( ST )
import Control.Monad.Primitive( PrimState
                              -- one day (GHC >= 7.10 ?)
                              , PrimMonad
                              )
import Data.Fixed( mod' )
import Graphics.Rasterific.Command
import Graphics.Rasterific.BiSampleable
import Graphics.Rasterific.Linear
             ( V2( .. )
             , (^-^)
             , (^/)
             , dot
             , norm
             )

import qualified Data.Vector as V

import Codec.Picture.Types( Pixel( .. )
                          , Image( .. )
                          , MutableImage( .. )
                          , Pixel8
                          , PixelRGBA8
                          , unsafeWritePixelBetweenAt
                          , readPackedPixelAt
                          , writePackedPixelAt
                          )

import Graphics.Rasterific.Types( Point
                                , Vector
                                , Line( .. )
                                , SamplerRepeat( .. ) )
import Graphics.Rasterific.Transformations
import Graphics.Rasterific.Rasterize
import Graphics.Rasterific.PatchTypes
import Graphics.Rasterific.Compositor( Modulable( .. )
                                     , ModulablePixel
                                     , RenderablePixel
                                     , compositionAlpha )


data TextureSpaceInfo = TextureSpaceInfo
    { _tsStart     :: {-# UNPACK #-} !Point
    , _tsDelta     :: {-# UNPACK #-} !Vector
    , _tsCoverage  :: {-# UNPACK #-} !Float
    , _tsRepeat    :: {-# UNPACK #-} !Int
    , _tsBaseIndex :: {-# UNPACK #-} !Int
    }
    deriving (Eq, Show)

type CoverageFiller m px =
    MutableImage (PrimState m) px -> CoverageSpan -> m ()

type Filler m =
    TextureSpaceInfo -> m ()

-- | Right now, we must stick to ST, due to the fact that
-- we can't specialize with parameterized monad :(
solidColor :: forall s px . (ModulablePixel px)
           => px -> MutableImage s px -> Filler (ST s)
{-# SPECIALIZE solidColor :: PixelRGBA8 -> MutableImage s PixelRGBA8
                          -> TextureSpaceInfo -> ST s () #-}
{-# SPECIALIZE solidColor :: Pixel8 -> MutableImage s Pixel8
                          -> TextureSpaceInfo -> ST s () #-}
solidColor color _ tsInfo
    | pixelOpacity color == emptyValue || _tsCoverage tsInfo <= 0 =
        return ()
solidColor color img tsInfo
    -- We are in the case fully opaque, so we can
    -- just overwrite what was there before
    | pixelOpacity color == fullOpacity && _tsCoverage tsInfo >= 1 =
        unsafeWritePixelBetweenAt img color (_tsBaseIndex tsInfo) maxi
        {-go 0 $ _tsBaseIndex tsInfo-}
  where
    !fullOpacity = fullValue :: PixelBaseComponent px
    !maxi = _tsRepeat tsInfo

-- We can be transparent, so perform alpha blending.
solidColor color img tsInfo = go 0 $ _tsBaseIndex tsInfo
  where
    !opacity = pixelOpacity color
    !(scanCoverage,_) = clampCoverage $_tsCoverage tsInfo
    !(cov, icov) = coverageModulate scanCoverage opacity
    !maxi = _tsRepeat tsInfo
    !compCount = componentCount (undefined :: px)

    go count  _ | count >= maxi = return ()
    go !count !idx = do
      oldPixel <- readPackedPixelAt img idx
      writePackedPixelAt img idx
        $ compositionAlpha cov icov oldPixel color
      go (count + 1) $ idx + compCount


-- | Plot a single pixel on the resulting image.
plotOpaquePixel :: forall m px. (ModulablePixel px, PrimMonad m)
                => MutableImage (PrimState m) px -> px -> Int -> Int
                -> m ()
{-# INLINE plotOpaquePixel #-}
plotOpaquePixel img _color x y
   | x < 0 || y < 0 || 
     x >= mutableImageWidth img || y >= mutableImageHeight img = return ()
plotOpaquePixel img color x y = do
  let !idx = (y * mutableImageWidth img + x) * (componentCount (undefined :: px))
  writePackedPixelAt img idx color

-- | Plot a single pixel on the resulting image, no bounds check are
-- performed, ensure index is correct!
unsafePlotOpaquePixel :: forall m px. (ModulablePixel px, PrimMonad m)
                      => MutableImage (PrimState m) px -> px -> Int -> Int
                      -> m ()
{-# INLINE unsafePlotOpaquePixel #-}
unsafePlotOpaquePixel img color x y = do
  let !idx = (y * mutableImageWidth img + x) * (componentCount (undefined :: px))
  writePackedPixelAt img idx color

shaderFiller :: forall s px . (ModulablePixel px)
             => ShaderFunction px -> MutableImage s px
             -> Filler (ST s)
{-# SPECIALIZE shaderFiller :: ShaderFunction PixelRGBA8
                            -> MutableImage s PixelRGBA8
                            -> Filler (ST s) #-}
{-# SPECIALIZE shaderFiller :: ShaderFunction Pixel8
                            -> MutableImage s Pixel8
                            -> Filler (ST s) #-}
shaderFiller shader img tsInfo =
    go 0 (_tsBaseIndex tsInfo) xStart yStart
  where
    !(scanCoverage,_) = clampCoverage $_tsCoverage tsInfo
    !maxi = _tsRepeat tsInfo
    !compCount = componentCount (undefined :: px)
    (V2 xStart yStart) = _tsStart tsInfo
    (V2 dx dy) = _tsDelta tsInfo

    go count  _ _ _ | count >= maxi = return ()
    go !count !idx !x !y = do
      let !color = shader x y
          !opacity = pixelOpacity color
          (cov, icov) = coverageModulate scanCoverage opacity
      oldPixel <- readPackedPixelAt img idx
      writePackedPixelAt img idx
        $ compositionAlpha cov icov oldPixel color
      go (count + 1) (idx + compCount) (x + dx) (y + dy)

prepareInfoNoTransform :: (Pixel px)
                       => MutableImage s px -> CoverageSpan
                       -> TextureSpaceInfo
prepareInfoNoTransform img coverage = TextureSpaceInfo
    { _tsStart     = V2 (_coverageX coverage) (_coverageY coverage)
    , _tsDelta     = V2 1 0
    , _tsCoverage  = _coverageVal coverage
    , _tsRepeat    = floor $ _coverageLength coverage
    , _tsBaseIndex =
        mutablePixelBaseIndex img (floor $ _coverageX coverage)
                                  (floor $ _coverageY coverage)
    }

prepareInfo :: (Pixel px)
            => Maybe Transformation -> MutableImage s px -> CoverageSpan
            -> TextureSpaceInfo
prepareInfo Nothing img covSpan = prepareInfoNoTransform img covSpan
prepareInfo (Just t) img covSpan = TextureSpaceInfo
    { _tsStart     = applyTransformation t
                   $ V2 (_coverageX covSpan) (_coverageY covSpan)
    , _tsDelta     = applyVectorTransformation t $ V2 1 0
    , _tsCoverage  = _coverageVal covSpan
    , _tsRepeat    = floor $ _coverageLength covSpan
    , _tsBaseIndex =
        mutablePixelBaseIndex img (floor $ _coverageX covSpan)
                                  (floor $ _coverageY covSpan)
    }

combineTransform :: Maybe Transformation -> Transformation
                 -> Maybe Transformation
combineTransform Nothing a = Just a
combineTransform (Just v) a = Just $ v <> a

withTrans :: Maybe Transformation -> ShaderFunction px
          -> ShaderFunction px
withTrans Nothing shader = shader
withTrans (Just v) shader = \x y ->
    let V2 x' y' = applyTransformation v (V2 x y) in
    shader x' y'

-- | The intent of shader texture is to provide ease of implementation
-- If possible providing a custom filler will be more efficient,
-- like already done for the solid colors.
shaderOfTexture :: forall px . RenderablePixel px
                => Maybe Transformation -> SamplerRepeat -> Texture px
                -> ShaderFunction px
{-# SPECIALIZE
    shaderOfTexture :: Maybe Transformation -> SamplerRepeat -> Texture PixelRGBA8
                    -> ShaderFunction PixelRGBA8 #-}
{-# SPECIALIZE
    shaderOfTexture :: Maybe Transformation -> SamplerRepeat -> Texture Pixel8
                    -> ShaderFunction Pixel8 #-}
shaderOfTexture _ _ (SolidTexture px) = \_ _ -> px
shaderOfTexture _ _ (MeshPatchTexture _ _) = error "MeshPatch should be precomputed"
shaderOfTexture trans sampling (LinearGradientTexture grad (Line a b)) =
  withTrans trans $ linearGradientShader grad a b sampling
shaderOfTexture trans sampling (RadialGradientTexture grad center radius) =
  withTrans trans $ radialGradientShader grad center radius sampling
shaderOfTexture trans sampling (RadialGradientWithFocusTexture grad center 
                                                    radius focus) =
  withTrans trans
             $ radialGradientWithFocusShader grad center radius focus
                                              sampling
shaderOfTexture trans _ (WithSampler sampler sub) =
  shaderOfTexture trans sampler sub
shaderOfTexture trans sampling (WithTextureTransform transform sub) =
  shaderOfTexture (combineTransform trans transform) sampling sub
shaderOfTexture trans sampling (SampledTexture img) =
  withTrans trans $ sampledImageShader img sampling
shaderOfTexture trans _ (ShaderTexture func) =
  withTrans trans func
shaderOfTexture trans _ (RawTexture img) =
  withTrans trans $ imageShader img
shaderOfTexture trans _sampling (PatternTexture _ _ _ _ img) =
  shaderOfTexture trans SamplerRepeat $ SampledTexture img
shaderOfTexture trans sampling (ModulateTexture texture modulation) =
  modulateTexture (shaderOfTexture trans sampling texture)
                  (shaderOfTexture trans sampling modulation)
shaderOfTexture trans sampling (AlphaModulateTexture texture modulation) =
  alphaModulateTexture
    (shaderOfTexture trans sampling texture)
    (shaderOfTexture trans sampling modulation)


-- | This function will interpret the texture description, helping
-- prepare and optimize the real calculation
transformTextureToFiller
    :: (RenderablePixel px)
    => (Maybe Transformation -> Int -> Int -> PatchInterpolation -> MeshPatch px -> Image px)
    -> Texture px -> CoverageFiller (ST s) px
transformTextureToFiller renderMesh = go Nothing SamplerPad
  where
    go _ _ (SolidTexture px) =
        \img -> solidColor px img . prepareInfoNoTransform img
    go trans sampling (WithTextureTransform transform sub) =
        go (combineTransform trans transform) sampling sub
    go trans _ (WithSampler sampler sub) =
        go trans sampler sub
    go trans sampling (MeshPatchTexture i m) = \img ->
      let newImg = renderMesh
            trans
            (mutableImageWidth img)
            (mutableImageHeight img)
            i
            m
      in
      go Nothing sampling (RawTexture newImg) img
        
    go trans sampling tex =
        \img -> shaderFiller shader img . prepareInfo trans img
            where shader = shaderOfTexture Nothing sampling tex

type GradientArray px = V.Vector (Float, px)

repeatGradient :: Float -> Float
repeatGradient s = s - fromIntegral (floor s :: Int)

reflectGradient :: Float -> Float
reflectGradient s =
    abs (abs (s - 1) `mod'` 2 - 1)
   
gradientColorAt :: ModulablePixel px
                => GradientArray px -> Float -> px
{-# SPECIALIZE
     gradientColorAt :: GradientArray PixelRGBA8 -> Float -> PixelRGBA8 #-}
{-# SPECIALIZE
     gradientColorAt :: GradientArray Pixel8 -> Float -> Pixel8 #-}
gradientColorAt grad at
    | at <= 0 = snd $ V.head grad
    | at >= 1.0 = snd $ V.last grad
    | otherwise = go (0, snd $ V.head grad) 0
  where
    !maxi = V.length grad
    go (prevCoeff, prevValue) ix
      | ix >= maxi = snd $ V.last grad
      | at < coeff = mixWith (\_ -> alphaOver cov icov) prevValue px
      | otherwise = go value $ ix + 1
      where value@(coeff, px) = grad `V.unsafeIndex` ix
            zeroToOne = (at - prevCoeff) / (coeff - prevCoeff)
            (cov, icov) = clampCoverage zeroToOne

gradientColorAtRepeat :: ModulablePixel px
                      => SamplerRepeat -> GradientArray px -> Float -> px
{-# SPECIALIZE INLINE
    gradientColorAtRepeat ::
        SamplerRepeat -> GradientArray PixelRGBA8 -> Float -> PixelRGBA8 #-}
{-# SPECIALIZE INLINE
    gradientColorAtRepeat ::
        SamplerRepeat -> GradientArray Pixel8 -> Float -> Pixel8 #-}
gradientColorAtRepeat SamplerPad grad = gradientColorAt grad
gradientColorAtRepeat SamplerRepeat grad =
    gradientColorAt grad . repeatGradient
gradientColorAtRepeat SamplerReflect grad =
    gradientColorAt grad . reflectGradient

linearGradientShader :: ModulablePixel px
                     => Gradient px -- ^ Gradient description.
                     -> Point       -- ^ Linear gradient start point.
                     -> Point       -- ^ Linear gradient end point.
                     -> SamplerRepeat
                     -> ShaderFunction px
{-# SPECIALIZE linearGradientShader
                     :: Gradient PixelRGBA8 -> Point -> Point -> SamplerRepeat
                     -> ShaderFunction PixelRGBA8 #-}
{-# SPECIALIZE linearGradientShader
                     :: Gradient Pixel8 -> Point -> Point -> SamplerRepeat
                     -> ShaderFunction Pixel8 #-}
linearGradientShader gradient start end repeating =
    \x y -> colorAt $ (V2 x y `dot` d) - s00
  where
    colorAt = gradientColorAtRepeat repeating gradArray
    gradArray = V.fromList gradient
    vector = end ^-^ start
    d = vector ^/ (vector `dot` vector)
    s00 = start `dot` d

-- | Use another image as a texture for the filling.
-- This texture use the "nearest" filtering, AKA no
-- filtering at all.
imageShader :: forall px. (Pixel px) => Image px -> ShaderFunction px
{-# SPECIALIZE
    imageShader :: Image PixelRGBA8 -> ShaderFunction PixelRGBA8 #-}
{-# SPECIALIZE
    imageShader :: Image Pixel8 -> ShaderFunction Pixel8 #-}
imageShader img x y =
    unsafePixelAt rawData $ (clampedY * w + clampedX) * compCount
  where
   clampedX = min (w - 1) . max 0 $ floor x
   clampedY = min (h - 1) . max 0 $ floor y
   !compCount = componentCount (undefined :: px)
   !w = imageWidth img
   !h = imageHeight img
   !rawData = imageData img

radialGradientShader :: ModulablePixel px
                     => Gradient px -- ^ Gradient description
                     -> Point       -- ^ Radial gradient center
                     -> Float       -- ^ Radial gradient radius
                     -> SamplerRepeat
                     -> ShaderFunction px
{-# SPECIALIZE 
    radialGradientShader
       :: Gradient PixelRGBA8 -> Point -> Float -> SamplerRepeat
       -> ShaderFunction PixelRGBA8 #-}
{-# SPECIALIZE 
    radialGradientShader
       :: Gradient Pixel8 -> Point -> Float -> SamplerRepeat
       -> ShaderFunction Pixel8 #-}
radialGradientShader gradient center radius repeating =
    \x y -> colorAt $ norm (V2 x y ^-^ center) / radius
  where
    !colorAt = gradientColorAtRepeat repeating gradArray
    !gradArray = V.fromList gradient

radialGradientWithFocusShader
    :: ModulablePixel px
    => Gradient px -- ^ Gradient description
    -> Point      -- ^ Radial gradient center
    -> Float      -- ^ Radial gradient radius
    -> Point      -- ^ Radial gradient focus point
    -> SamplerRepeat
    -> ShaderFunction px
{-# SPECIALIZE
    radialGradientWithFocusShader
        :: Gradient PixelRGBA8 -> Point -> Float -> Point
        -> SamplerRepeat -> ShaderFunction PixelRGBA8 #-}
{-# SPECIALIZE
    radialGradientWithFocusShader
        :: Gradient Pixel8 -> Point -> Float -> Point
        -> SamplerRepeat -> ShaderFunction Pixel8 #-}
radialGradientWithFocusShader gradient center radius focusScreen repeating =
    \x y -> colorAt . go $ V2 x y ^-^ center
  where
    focus@(V2 origFocusX origFocusY) = focusScreen ^-^ center
    colorAt = gradientColorAtRepeat repeating gradArray
    gradArray = V.fromList gradient
    radiusSquared = radius * radius
    dist = sqrt $ focus `dot` focus
    clampedFocus@(V2 focusX focusY)
        | dist <= r = focus
        | otherwise = V2 (r * cos a) (r * sin a)
           where a = atan2 origFocusY origFocusX
                 r = radius * 0.99
    trivial = sqrt $ radiusSquared - origFocusX * origFocusX

    solutionOf (V2 x y) | x == focusX =
        V2 focusX (if y > focusY then trivial else negate trivial)
    solutionOf (V2 x y) = V2 xSolution $ slope * xSolution + yint
      where
        slope = (y - focusY) / (x - focusX)
        yint = y - (slope * x)

        a = slope * slope + 1
        b = 2 * slope * yint
        c = yint * yint - radiusSquared
        det = sqrt $ b * b - 4 * a * c
        xSolution = (-b + (if x < focusX then negate det else det)) / (2 * a)

    go pos = sqrt $ curToFocus / distSquared
      where
        solution = solutionOf pos ^-^ clampedFocus
        toFocus = pos ^-^ clampedFocus
        distSquared = solution `dot` solution
        curToFocus = toFocus `dot` toFocus

-- | Perform a multiplication operation between a full color texture
-- and a greyscale one, used for clip-path implementation.
modulateTexture :: ModulablePixel px
                => ShaderFunction px
                -> ShaderFunction (PixelBaseComponent px)
                -> ShaderFunction px
{-# INLINE modulateTexture #-}
modulateTexture fullTexture modulator x y =
    colorMap (modulate $ modulator x y) $ fullTexture x y

-- | Perform a multiplication operation between a full color texture
-- and a greyscale one, used for clip-path implementation.
alphaModulateTexture :: ModulablePixel px
                => ShaderFunction px
                -> ShaderFunction (PixelBaseComponent px)
                -> ShaderFunction px
{-# INLINE alphaModulateTexture #-}
alphaModulateTexture fullTexture modulator x y =
  let px = fullTexture x y in
  mixWithAlpha (\_ _ a -> a) (\_ _ -> modulator x y) px px