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

module Spec.Quaternion (tests) where

import Geomancy
import Hedgehog
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testPropertyNamed)

import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Linear.Quaternion qualified

import Geomancy.Quaternion qualified
import Geomancy.Transform ((!.))
import Geomancy.Transform qualified as Transform

import Support.Ball (balls, consistent, enclosed, exactly)
import Support.Fixtures
  (intrinsicComposed, intrinsicComposedRef, quarterTurn, rotationsMatch, xNeg, xPos, yNeg, yPos, zNeg, zPos)
import Support.Gen (genAngle, genSmallVec3, genVec3)
import Support.Linear (toQ, toV3)
import Support.Reference.Linear qualified as Ref.Linear
import Support.Reference.Quaternion qualified as Ref.Quaternion
import Support.Reference.Transform qualified as Ref.Transform

tests :: TestTree
tests = testGroup "Quaternion"
  [ testPropertyNamed "ref" "prop_quaternion_ref" prop_quaternion_ref
  , testPropertyNamed "rotate_matrix" "prop_quaternion_rotate_matrix" prop_quaternion_rotate_matrix
  , testPropertyNamed "associativity" "prop_quaternion_associativity" prop_quaternion_associativity
  , testPropertyNamed "intrinsic" "prop_quaternion_intrinsic" prop_quaternion_intrinsic
  , testPropertyNamed "extrinsic" "prop_quaternion_extrinsic" prop_quaternion_extrinsic
  , testPropertyNamed "intrinsic_ref" "prop_quaternion_intrinsic_ref" prop_quaternion_intrinsic_ref
  ]

prop_quaternion_ref :: Property
prop_quaternion_ref = property do
  annotate "Construction"
  (axis, angle) <- forAll $ (,) <$> genVec3 <*> genAngle
  let qg = Geomancy.Quaternion.axisAngle axis angle
  let ql = Linear.Quaternion.axisAngle (toV3 axis) angle
  show (toQ qg) === show ql

  annotate "Rotation"
  p <- forAll genSmallVec3
  let
    rpg = Geomancy.Quaternion.rotate qg p
    rpl = Linear.Quaternion.rotate ql (toV3 p)
    rpgBall =
      Ref.Quaternion.rotate
        (Ref.Quaternion.axisAngle (balls (toV3 axis)) (exactly angle))
        (balls (toV3 p))
    rplBall =
      Ref.Linear.rotate
        (Linear.Quaternion.axisAngle (balls (toV3 axis)) (exactly angle))
        (balls (toV3 p))
  replicate 3 Nothing === enclosed rpgBall (toV3 rpg)
  replicate 3 Nothing === enclosed rplBall rpl
  replicate 3 Nothing === consistent rpgBall rplBall

  annotate "Multiplication"
  (axisB, angleB) <- forAll $ (,) <$> genVec3 <*> genAngle
  let qgB = Geomancy.Quaternion.axisAngle axisB angleB
  let qlB = Linear.Quaternion.axisAngle (toV3 axisB) angleB
  show (toQ (qg * qgB)) === show (ql * qlB)

prop_quaternion_rotate_matrix :: Property
prop_quaternion_rotate_matrix = property do
  q <- forAll $ quaternion <$> genUnitish <*> genUnitish <*> genUnitish <*> genUnitish
  p <- forAll genSmallVec3
  let
    viaRotate = Geomancy.Quaternion.rotate q p
    viaMatrix = Transform.rotateQ q !. p
    viaRotateBall = Ref.Quaternion.rotate (balls (toQ q)) (balls (toV3 p))
    viaMatrixBall = Ref.Transform.rotateQ (balls (toQ q)) Ref.Transform.!. balls (toV3 p)
  replicate 3 Nothing === enclosed viaRotateBall (toV3 viaRotate)
  replicate 3 Nothing === enclosed viaMatrixBall (toV3 viaMatrix)
  replicate 3 Nothing === consistent viaRotateBall viaMatrixBall
  where
    genUnitish = Gen.float (Range.linearFracFrom 0.0 (-1) 1)

prop_quaternion_associativity :: Property
prop_quaternion_associativity = property do
  (a, b, c) <- forAll $ (,,)
    <$> (Geomancy.Quaternion.axisAngle <$> genVec3 <*> genAngle)
    <*> (Geomancy.Quaternion.axisAngle <$> genVec3 <*> genAngle)
    <*> (Geomancy.Quaternion.axisAngle <$> genVec3 <*> genAngle)
  let ab'c = (a * b) * c
  annotateShow ab'c
  let a'bc = a * (b * c)
  annotateShow a'bc
  let
    aBall = balls (toQ a)
    bBall = balls (toQ b)
    cBall = balls (toQ c)
    ab'cBall = (aBall `Ref.Quaternion.multiply` bBall) `Ref.Quaternion.multiply` cBall
    a'bcBall = aBall `Ref.Quaternion.multiply` (bBall `Ref.Quaternion.multiply` cBall)
  replicate 4 Nothing === enclosed ab'cBall (toQ ab'c)
  replicate 4 Nothing === enclosed a'bcBall (toQ a'bc)
  replicate 4 Nothing === consistent ab'cBall a'bcBall

prop_quaternion_intrinsic :: Property
prop_quaternion_intrinsic = withTests 1 $ property do
  let
    q = Geomancy.Quaternion.intrinsic 0 quarterTurn quarterTurn
    qBall = Ref.Quaternion.intrinsic 0 quarterTurn quarterTurn
  rotationsMatch toV3
    (Geomancy.Quaternion.rotate q)
    (Ref.Quaternion.rotate qBall)
    [(zPos, yNeg), (xPos, zNeg), (yNeg, xNeg)]

prop_quaternion_extrinsic :: Property
prop_quaternion_extrinsic = withTests 1 $ property do
  let
    q = Geomancy.Quaternion.extrinsic quarterTurn quarterTurn 0
    qBall = Ref.Quaternion.extrinsic quarterTurn quarterTurn 0
  rotationsMatch toV3
    (Geomancy.Quaternion.rotate q)
    (Ref.Quaternion.rotate qBall)
    [(zPos, xPos), (xPos, yPos), (yNeg, zNeg)]

prop_quaternion_intrinsic_ref :: Property
prop_quaternion_intrinsic_ref = property do
  (roll, pitch, yaw) <- forAll $ (,,) <$> genAngle <*> genAngle <*> genAngle
  let fused = Geomancy.Quaternion.intrinsic roll pitch yaw
  annotateShow fused
  let composed = intrinsicComposed roll pitch yaw
  annotateShow composed
  let viaFused = Geomancy.Quaternion.rotate fused zPos
  annotateShow viaFused
  let viaComposed = Geomancy.Quaternion.rotate composed zPos
  let
    zPosBall = balls (toV3 zPos)
    viaFusedBall =
      Ref.Quaternion.rotate (Ref.Quaternion.intrinsic (exactly roll) (exactly pitch) (exactly yaw)) zPosBall
    viaComposedBall =
      Ref.Quaternion.rotate (intrinsicComposedRef (exactly roll) (exactly pitch) (exactly yaw)) zPosBall
  replicate 3 Nothing === enclosed viaFusedBall (toV3 viaFused)
  replicate 3 Nothing === enclosed viaComposedBall (toV3 viaComposed)
  replicate 3 Nothing === consistent viaFusedBall viaComposedBall