{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DerivingStrategies #-}

-- | Specialized and inlined @V3 Float@.

module Geomancy.Vec3
  ( Vec3
  , vec3
  , withVec3
  , pattern WithVec3
  , convert
  , fromVec2
  , fromTuple

  , (^*)
  , (^/)
  , lerp

  , cross
  , dot
  , normalize

  , Packed(..)
  , packed

  , emap2
  , emap3
  , emap4
  ) where

import Control.DeepSeq (NFData(rnf))
import Data.Coerce (Coercible, coerce)
import Data.MonoTraversable (Element, MonoFunctor(..), MonoPointed(..))
import Data.VectorSpace (VectorSpace)
import Foreign (Storable(..), castPtr)
import Foreign.Ptr.Diff (peekDiffOff, pokeDiffOff)
import qualified Data.VectorSpace as VectorSpace

import Geomancy.Elementwise (Elementwise(..))
import Graphics.Gl.Block (Block(..))
import Geomancy.Gl.Funs (GlModf(..), GlNearest)
import Geomancy.Vec2 (Vec2, withVec2)

data Vec3 = Vec3
  {-# UNPACK #-} !Float
  {-# UNPACK #-} !Float
  {-# UNPACK #-} !Float
  deriving (Eq, Ord, Show)

{-# INLINE vec3 #-}
vec3 :: Float -> Float -> Float -> Vec3
vec3 = Vec3

{-# INLINE withVec3 #-}
withVec3
  :: Vec3
  -> (Float -> Float -> Float -> r)
  -> r
withVec3 (Vec3 a b c) f = f a b c

{-# INLINE convert #-}
convert :: Coercible v Vec3 => (Float -> a) -> (a -> a -> a -> r) -> v -> r
convert f t v =
  withVec3 (coerce v) \a b c ->
    t (f a) (f b) (f c)

pattern WithVec3 :: Float -> Float -> Float -> Vec3
pattern WithVec3 a b c <- ((`withVec3` (,,)) -> (a, b, c))
{-# COMPLETE WithVec3 #-}

{-# INLINE fromVec2 #-}
fromVec2 :: Coercible Vec3 a => Vec2 -> Float -> a
fromVec2 xy z =
  withVec2 xy \x y ->
    coerce (vec3 x y z)

{-# INLINE fromTuple #-}
fromTuple :: Coercible Vec3 a => (Float, Float, Float) -> a
fromTuple (x, y, z) = coerce (vec3 x y z)

instance NFData Vec3 where
  rnf Vec3{} = ()

type instance Element Vec3 = Float

instance MonoFunctor Vec3 where
  {-# INLINE omap #-}
  omap f v =
    withVec3 v \x y z ->
      vec3 (f x) (f y) (f z)

instance MonoPointed Vec3 where
  opoint x = vec3 x x x

instance Elementwise Vec3 where
  {-# INLINE emap2 #-}
  emap2 f p0 p1 =
    withVec3 p0 \x0 y0 z0 ->
    withVec3 p1 \x1 y1 z1 ->
      vec3
        (f x0 x1)
        (f y0 y1)
        (f z0 z1)

  {-# INLINE emap3 #-}
  emap3 f p0 p1 p2 =
    withVec3 p0 \x0 y0 z0 ->
    withVec3 p1 \x1 y1 z1 ->
    withVec3 p2 \x2 y2 z2 ->
      vec3
        (f x0 x1 x2)
        (f y0 y1 y2)
        (f z0 z1 z2)

  {-# INLINE emap4 #-}
  emap4 f p0 p1 p2 p3 =
    withVec3 p0 \x0 y0 z0 ->
    withVec3 p1 \x1 y1 z1 ->
    withVec3 p2 \x2 y2 z2 ->
    withVec3 p3 \x3 y3 z3 ->
      vec3
        (f x0 x1 x2 x3)
        (f y0 y1 y2 y3)
        (f z0 z1 z2 z3)

  {-# INLINE emap5 #-}
  emap5 f p0 p1 p2 p3 p4 =
    withVec3 p0 \x0 y0 z0 ->
    withVec3 p1 \x1 y1 z1 ->
    withVec3 p2 \x2 y2 z2 ->
    withVec3 p3 \x3 y3 z3 ->
    withVec3 p4 \x4 y4 z4 ->
      vec3
        (f x0 x1 x2 x3 x4)
        (f y0 y1 y2 y3 y4)
        (f z0 z1 z2 z3 z4)

instance Num Vec3 where
  {-# INLINE (+) #-}
  Vec3 a b c + Vec3 d e f =
    Vec3
      (a + d)
      (b + e)
      (c + f)

  {-# INLINE (-) #-}
  Vec3 a b c - Vec3 d e f =
    Vec3
      (a - d)
      (b - e)
      (c - f)

  {-# INLINE (*) #-}
  Vec3 a b c * Vec3 d e f =
    Vec3
      (a * d)
      (b * e)
      (c * f)

  {-# INLINE abs #-}
  abs (Vec3 a b c) =
    Vec3 (abs a) (abs b) (abs c)

  {-# INLINE signum #-}
  signum (Vec3 a b c) =
    Vec3 (signum a) (signum b) (signum c)

  {-# INLINE fromInteger #-}
  fromInteger x = Vec3 x' x' x'
    where
      x' = fromInteger x

instance Fractional Vec3 where
  {-# INLINE (/) #-}
  Vec3 l1 l2 l3 / Vec3 r1 r2 r3 =
    Vec3 (l1 / r1) (l2 / r2) (l3 / r3)

  {-# INLINE recip #-}
  recip (Vec3 a b c) =
    Vec3 (recip a) (recip b) (recip c)

  {-# INLINE fromRational #-}
  fromRational x = Vec3 x' x' x'
    where
      x' = fromRational x

instance Floating Vec3 where
  pi = opoint pi

  exp = omap exp
  log = omap log
  sqrt = omap sqrt
  sin = omap sin
  cos = omap cos
  asin = omap asin
  acos = omap acos
  atan = omap atan
  sinh = omap sinh
  cosh = omap cosh
  asinh = omap asinh
  acosh = omap acosh
  atanh = omap atanh

  a ** b =
    withVec3 a \ax ay az ->
    withVec3 b \bx by bz ->
      vec3
        (ax ** bx)
        (ay ** by)
        (az ** bz)

{-
  XXX: GPU layouts call for some padding.

  Maybe it would be worth it to flip the sizeOf-s.
-}
instance Storable Vec3 where
  {-# INLINE sizeOf #-}
  sizeOf _ = 16

  {-# INLINE alignment #-}
  alignment _ = 4

  {-# INLINE poke #-}
  poke ptr v3 =
    withVec3 v3 \a b c -> do
      poke ptr' a
      pokeElemOff ptr' 1 b
      pokeElemOff ptr' 2 c
      pokeElemOff ptr' 3 (1.0 :: Float)
    where
      ptr' = castPtr ptr

  {-# INLINE peek #-}
  peek ptr =
    vec3 <$> peek ptr' <*> peekElemOff ptr' 1 <*> peekElemOff ptr' 2
    where
      ptr' = castPtr ptr

instance Block Vec3 where
  type PackedSize Vec3 = 12
  alignment140 _  = 16
  sizeOf140 _     = 16
  alignment430    = alignment140
  sizeOf430       = sizeOf140
  isStruct _      = False
  read140     = peekDiffOff
  write140    = pokeDiffOff
  read430     = read140
  write430    = write140
  readPacked  = read140
  writePacked = write140
  {-# INLINE alignment140 #-}
  {-# INLINE sizeOf140 #-}
  {-# INLINE alignment430 #-}
  {-# INLINE sizeOf430 #-}
  {-# INLINE isStruct #-}
  {-# INLINE read140 #-}
  {-# INLINE write140 #-}
  {-# INLINE read430 #-}
  {-# INLINE write430 #-}
  {-# INLINE readPacked #-}
  {-# INLINE writePacked #-}

{-# INLINE (^*) #-}
(^*) :: Vec3 -> Float -> Vec3
Vec3 a b c ^* x =
  Vec3
    (a * x)
    (b * x)
    (c * x)

{-# INLINE (^/) #-}
(^/) :: Vec3 -> Float -> Vec3
Vec3 a b c ^/ x =
  Vec3
    (a / x)
    (b / x)
    (c / x)

{-# INLINE lerp #-}
lerp :: Float -> Vec3 -> Vec3 -> Vec3
lerp alpha u v = u ^* alpha + v ^* (1 - alpha)

{-# INLINE cross #-}
cross :: Vec3 -> Vec3 -> Vec3
cross (Vec3 a b c) (Vec3 d e f) =
  Vec3
    (b * f - c * e)
    (c * d - a * f)
    (a * e - b * d)

{-# INLINE dot #-}
dot :: Vec3 -> Vec3 -> Float
dot (Vec3 a b c) (Vec3 d e f) =
  a * d +
  b * e +
  c * f

{-# INLINE normalize #-}
normalize :: Vec3 -> Vec3
normalize v =
  if nearZero q || nearZero (1-q) then
    v
  else
    let
      Vec3 x y z = v
    in
      Vec3 (x / l) (y / l) (z / l)

  where
    q = dot v v
    l = sqrt q

    nearZero a = abs a <= 1e-6

instance VectorSpace Vec3 Float where
  zeroVector = epoint 0

  {-# INLINE (*^) #-}
  (*^) = flip (Geomancy.Vec3.^*)

  {-# INLINE (^/) #-}
  (^/) = (Geomancy.Vec3.^/)

  {-# INLINE (^+^) #-}
  (^+^) = emap2 (+)

  {-# INLINE (^-^) #-}
  (^-^) = emap2 (-)

  {-# INLINE negateVector #-}
  negateVector = emap negate

  {-# INLINE dot #-}
  dot = Geomancy.Vec3.dot

  {-# INLINE normalize #-}
  normalize = Geomancy.Vec3.normalize

-- * Unpadded

type instance Element Packed = Float

newtype Packed = Packed { unPacked :: Vec3 }
  deriving stock
    ( Eq, Ord, Show
    )
  deriving newtype
    ( NFData, Num, Fractional, Floating
    , MonoFunctor, MonoPointed
    , Elementwise
    )

{-# INLINE packed #-}
packed :: Float -> Float -> Float -> Packed
packed x y z = Packed (vec3 x y z)

instance Storable Packed where
  {-# INLINE sizeOf #-}
  sizeOf _ = 12

  {-# INLINE alignment #-}
  alignment _ = 4

  {-# INLINE poke #-}
  poke ptr (Packed v3) =
    withVec3 v3 \a b c -> do
      poke ptr' a
      pokeElemOff ptr' 1 b
      pokeElemOff ptr' 2 c
    where
      ptr' = castPtr ptr

  {-# INLINE peek #-}
  peek ptr = packed
    <$> peek ptr'
    <*> peekElemOff ptr' 1
    <*> peekElemOff ptr' 2
    where
      ptr' = castPtr ptr

instance Block Packed where
  type PackedSize Packed = 12
  alignment140 _  = 16
  sizeOf140 _     = 16
  alignment430    = alignment140
  sizeOf430       = sizeOf140
  isStruct _      = False
  read140     = peekDiffOff
  write140    = pokeDiffOff
  read430     = read140
  write430    = write140
  readPacked  = read140
  writePacked = write140
  {-# INLINE alignment140 #-}
  {-# INLINE sizeOf140 #-}
  {-# INLINE alignment430 #-}
  {-# INLINE sizeOf430 #-}
  {-# INLINE isStruct #-}
  {-# INLINE read140 #-}
  {-# INLINE write140 #-}
  {-# INLINE read430 #-}
  {-# INLINE write430 #-}
  {-# INLINE readPacked #-}
  {-# INLINE writePacked #-}

instance VectorSpace Packed Float where
  zeroVector = epoint 0

  {-# INLINE (*^) #-}
  (*^) = flip $ coerce (Geomancy.Vec3.^*)

  {-# INLINE (^/) #-}
  (^/) = coerce (Geomancy.Vec3.^/)

  {-# INLINE (^+^) #-}
  v1 ^+^ v2 = v1 + v2

  {-# INLINE (^-^) #-}
  v1 ^-^ v2 = v1 - v2

  {-# INLINE negateVector #-}
  negateVector = emap negate

  {-# INLINE dot #-}
  dot = coerce Geomancy.Vec3.dot

  {-# INLINE normalize #-}
  normalize = coerce Geomancy.Vec3.normalize

instance GlNearest Vec3

instance GlModf Vec3 Vec3 where
  glModf v =
    withVec3 v \vx vy vz ->
      let
        (xi, xf) = glModf vx
        (yi, yf) = glModf vy
        (zi, zf) = glModf vz
      in
        ( vec3 (fromInteger xi) (fromInteger yi) (fromInteger zi)
        , vec3 xf yf zf
        )