{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Support.Swizzle
  ( Lanes(..)
  , Mismatch
  , swizzled
  , swizzles
  ) where

import Data.Int (Int32, Int64)
import Data.Maybe (catMaybes, fromMaybe)
import Data.Proxy (Proxy(..))
import Data.Word (Word32, Word64)
import GHC.Float (castDoubleToWord64, castFloatToWord32)
import GHC.Records (HasField(..))
import GHC.TypeLits (KnownSymbol, symbolVal)
import Hedgehog

import Geomancy

import Support.Property (pattern PROP_TESTS)

class Lanes v where
  lanes :: v -> [Word64]

instance Lanes Float where
  lanes x = [fromIntegral (castFloatToWord32 x)]

instance Lanes Double where
  lanes x = [castDoubleToWord64 x]

instance Lanes Int32 where
  lanes x = [fromIntegral x]

instance Lanes Word32 where
  lanes x = [fromIntegral x]

instance Lanes Int64 where
  lanes x = [fromIntegral x]

instance Lanes Word64 where
  lanes x = [x]

instance Lanes Vec2 where
  lanes v = withVec2 v \x y -> concatMap lanes [x, y]

instance Lanes Vec3 where
  lanes v = withVec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes Vec4 where
  lanes v = withVec4 v \x y z w -> concatMap lanes [x, y, z, w]

instance Lanes DVec2 where
  lanes v = withDVec2 v \x y -> concatMap lanes [x, y]

instance Lanes DVec3 where
  lanes v = withDVec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes DVec4 where
  lanes v = withDVec4 v \x y z w -> concatMap lanes [x, y, z, w]

instance Lanes IVec2 where
  lanes v = withIVec2 v \x y -> concatMap lanes [x, y]

instance Lanes IVec3 where
  lanes v = withIVec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes IVec4 where
  lanes v = withIVec4 v \x y z w -> concatMap lanes [x, y, z, w]

instance Lanes UVec2 where
  lanes v = withUVec2 v \x y -> concatMap lanes [x, y]

instance Lanes UVec3 where
  lanes v = withUVec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes UVec4 where
  lanes v = withUVec4 v \x y z w -> concatMap lanes [x, y, z, w]

instance Lanes I64Vec2 where
  lanes v = withI64Vec2 v \x y -> concatMap lanes [x, y]

instance Lanes I64Vec3 where
  lanes v = withI64Vec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes I64Vec4 where
  lanes v = withI64Vec4 v \x y z w -> concatMap lanes [x, y, z, w]

instance Lanes U64Vec2 where
  lanes v = withU64Vec2 v \x y -> concatMap lanes [x, y]

instance Lanes U64Vec3 where
  lanes v = withU64Vec3 v \x y z -> concatMap lanes [x, y, z]

instance Lanes U64Vec4 where
  lanes v = withU64Vec4 v \x y z w -> concatMap lanes [x, y, z, w]

data Mismatch = Mismatch String [Word64] [Word64]
  deriving (Eq, Show)

swizzles :: Show v => Gen v -> (v -> [Maybe Mismatch]) -> Property
swizzles gen checks = withTests PROP_TESTS $ property do
  v <- forAll gen
  catMaybes (checks v) === []

swizzled
  :: forall s v r
  .  (KnownSymbol s, HasField s v r, Lanes v, Lanes r, Num r)
  => v
  -> Maybe Mismatch
swizzled v =
  if expected == actual then
    Nothing
  else
    Just (Mismatch name expected actual)
  where
    name = symbolVal (Proxy @s)
    expected = concatMap lane name
    actual = lanes (getField @s v)
    lane c = fromMaybe (error $ "no lane " <> [c]) $ lookup c (zip "xyzw" (map pure (lanes v)) <> constants)
    constants = [('0', splat 0), ('1', splat 1)]
    splat x = take 1 (lanes (x :: r))