{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ImportQualifiedPost #-}

module Support.Gen
  ( genFloat
  , genDouble
  , genAngle
  , genAngleD
  , genVec2
  , genVec3
  , genVec4
  , genSmallVec3
  , genDVec2
  , genDVec3
  , genDVec4
  , genSmallDVec3
  , genDQuaternion
  , genInt32
  , genWord32
  , genInt64
  , genWord64
  , genEdgeScalar
  , genMostlySame
  , genTransform
  , forAllTransform
  , forAllDTransform
  ) where

import Geomancy
import Hedgehog

import Control.Monad (when)
import Data.Int (Int32, Int64)
import Data.Word (Word32, Word64)
import GHC.Stack (withFrozenCallStack)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Linear qualified
import Text.Printf (printf)

import Geomancy.DMat4 qualified
import Geomancy.Mat4 qualified
import Geomancy.Transform qualified as Transform

import Support.Linear (toM44, toM44D)

forAllTransform :: PropertyT IO (Geomancy.Mat4.Mat4, Linear.M44 Float)
forAllTransform = withFrozenCallStack do
  (_name, Transform g) <- forAllWith fst genTransform
  pure (g, toM44 g)

forAllDTransform :: PropertyT IO (Geomancy.DMat4.DMat4, Linear.M44 Double)
forAllDTransform = withFrozenCallStack do
  (_name, t) <- forAllWith fst genTransform
  let g = Geomancy.DMat4.fromMat4 t
  pure (g, toM44D g)

genInt32 :: Gen Int32
genInt32 = Gen.int32 Range.linearBounded

genWord32 :: Gen Word32
genWord32 = Gen.word32 Range.linearBounded

genInt64 :: Gen Int64
genInt64 = Gen.int64 Range.linearBounded

genWord64 :: Gen Word64
genWord64 = Gen.word64 Range.linearBounded

genEdgeScalar :: Fractional a => Gen a
genEdgeScalar = Gen.frequency [(4, pure 0), (4, pure (-0)), (4, pure 1), (4, pure (-1)), (1, pure (0 / 0))]

genMostlySame :: Fractional a => Int -> a -> Gen a
genMostlySame sameWeight x = Gen.frequency [(sameWeight, pure x), (1, genEdgeScalar)]

genFloat :: Gen Float
genFloat = Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)

genDouble :: Gen Double
genDouble = Gen.double (Range.linearFracFrom 0.0 (-1e6) 1e6)

genVec2 :: Gen Vec2
genVec2 = vec2 <$> genFloat <*> genFloat

genVec4 :: Gen Vec4
genVec4 = vec4 <$> genFloat <*> genFloat <*> genFloat <*> genFloat

genDVec2 :: Gen DVec2
genDVec2 = dvec2 <$> genDouble <*> genDouble

genDVec3 :: Gen DVec3
genDVec3 = dvec3 <$> genDouble <*> genDouble <*> genDouble

genDVec4 :: Gen DVec4
genDVec4 = dvec4 <$> genDouble <*> genDouble <*> genDouble <*> genDouble

genSmallVec3 :: Gen Vec3
genSmallVec3 =
  vec3
    <$> Gen.float (Range.linearFracFrom 0.0 (-10) 10)
    <*> Gen.float (Range.linearFracFrom 0.0 (-10) 10)
    <*> Gen.float (Range.linearFracFrom 0.0 (-10) 10)

genSmallDVec3 :: Gen DVec3
genSmallDVec3 =
  dvec3
    <$> Gen.double (Range.linearFracFrom 0.0 (-10) 10)
    <*> Gen.double (Range.linearFracFrom 0.0 (-10) 10)
    <*> Gen.double (Range.linearFracFrom 0.0 (-10) 10)

genDQuaternion :: Gen DQuaternion
genDQuaternion = dquaternion <$> genDouble <*> genDouble <*> genDouble <*> genDouble

genAngleD :: Gen Double
genAngleD = Gen.double (Range.linearFracFrom 0.0 (-8*pi) (8*pi))

genVec3 :: Gen Vec3
genVec3 =
  vec3
    <$> Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)
    <*> Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)
    <*> Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)

genAngle :: Gen Float
genAngle = Gen.float (Range.linearFracFrom 0.0 (-8*pi) (8*pi))

genTransform :: Gen ([Char], Transform) -- TODO: cross-check with Linear
genTransform = Gen.choice
  [ genIdentity
  , genTranslate
  , genRotate
  , genScale
  ]
  where
    genIdentity = pure ("identity", mempty)

    genTranslate = do
      x <- Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)
      y <- Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)
      z <- Gen.float (Range.linearFracFrom 0.0 (-1e6) 1e6)
      when (abs x == 0 && abs y == 0 && abs z == 0) Gen.discard
      pure
        ( printf "translate %0.4f %0.4f %0.4f" x y z
        , Transform.translate x y z
        )

    genRotate = do
      (name, axis) <- Gen.element
        [ ("rotate/x", Transform.rotateX)
        , ("rotate/y", Transform.rotateY)
        , ("rotate/z", Transform.rotateZ)
        ]
      angle <- genAngle
      when (abs angle == 0) Gen.discard
      pure
        ( printf "%s %0.4f" name angle
        , axis angle
        )

    genScale = do
      x <- Gen.float (Range.linearFracFrom 1.0 1e-6 1e6)
      y <- Gen.float (Range.linearFracFrom 1.0 1e-6 1e6)
      z <- Gen.float (Range.linearFracFrom 1.0 1e-6 1e6)
      when (abs x == 1 && abs y == 1 && abs z == 1) Gen.discard
      pure
        ( printf "scale %0.4f %0.4f %0.4f" x y z
        , Transform.scale3 x y z
        )