Quaternion.hs

module Rattletrap.Type.Quaternion where

import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Data.Ord as Ord
import qualified Rattletrap.BitGet as BitGet
import qualified Rattletrap.BitPut as BitPut
import qualified Rattletrap.Exception.InvalidComponent as InvalidComponent
import qualified Rattletrap.Schema as Schema
import qualified Rattletrap.Type.CompressedWord as CompressedWord
import qualified Rattletrap.Utility.Json as Json

data Quaternion = Quaternion
  { x :: Double,
    y :: Double,
    z :: Double,
    w :: Double
  }
  deriving (Eq, Show)

instance Json.FromJSON Quaternion where
  parseJSON = Json.withObject "Quaternion" $ \object -> do
    x <- Json.required object "x"
    y <- Json.required object "y"
    z <- Json.required object "z"
    w <- Json.required object "w"
    pure Quaternion {x, y, z, w}

instance Json.ToJSON Quaternion where
  toJSON a =
    Json.object
      [ Json.pair "x" $ x a,
        Json.pair "y" $ y a,
        Json.pair "z" $ z a,
        Json.pair "w" $ w a
      ]

schema :: Schema.Schema
schema =
  Schema.named "quaternion" $
    Schema.object
      [ (Json.pair "x" $ Schema.ref Schema.number, True),
        (Json.pair "y" $ Schema.ref Schema.number, True),
        (Json.pair "z" $ Schema.ref Schema.number, True),
        (Json.pair "w" $ Schema.ref Schema.number, True)
      ]

data Component
  = X
  | Y
  | Z
  | W
  deriving (Eq, Show)

toQuaternion :: Component -> Double -> Double -> Double -> Quaternion
toQuaternion component a b c =
  let d = toPart a b c
   in case component of
        X -> Quaternion d a b c
        Y -> Quaternion a d b c
        Z -> Quaternion a b d c
        W -> Quaternion a b c d

toPart :: Double -> Double -> Double -> Double
toPart a b c = sqrt (1 - (a * a) - (b * b) - (c * c))

compressPart :: Double -> CompressedWord.CompressedWord
compressPart =
  CompressedWord.CompressedWord maxCompressedValue
    . round
    . (* wordToDouble maxCompressedValue)
    . (+ 0.5)
    . (/ 2.0)
    . (/ maxValue)

decompressPart :: CompressedWord.CompressedWord -> Double
decompressPart x_ =
  (* maxValue)
    . (* 2.0)
    . subtract 0.5
    . (/ wordToDouble (CompressedWord.limit x_))
    . wordToDouble
    $ CompressedWord.value x_

maxComponent :: Quaternion -> Component
maxComponent quaternion =
  let x_ = x quaternion
      y_ = y quaternion
      z_ = z quaternion
      w_ = w quaternion
      parts = [(x_, X), (y_, Y), (z_, Z), (w_, W)]
      biggestPart = maximumOn fst parts
      roundTrip = decompressPart . compressPart
      computedPart =
        Maybe.fromMaybe
          biggestPart
          (List.find (\(value, _) -> value /= roundTrip value) parts)
   in snd
        ( if (biggestPart == computedPart)
            || (abs (fst biggestPart - fst computedPart) > 0.00001)
            then biggestPart
            else computedPart
        )

maximumOn :: (Foldable t, Ord b) => (a -> b) -> t a -> a
maximumOn f = List.maximumBy (Ord.comparing f)

numBits :: Word
numBits = 18

wordToDouble :: Word -> Double
wordToDouble = fromIntegral

maxCompressedValue :: Word
maxCompressedValue = (2 ^ numBits) - 1

maxValue :: Double
maxValue = 1.0 / sqrt 2.0

bitPut :: Quaternion -> BitPut.BitPut
bitPut q =
  let c = maxComponent q
   in putComponent c <> case c of
        X -> putParts (y q) (z q) (w q)
        Y -> putParts (x q) (z q) (w q)
        Z -> putParts (x q) (y q) (w q)
        W -> putParts (x q) (y q) (z q)

putComponent :: Component -> BitPut.BitPut
putComponent component =
  CompressedWord.bitPut
    ( CompressedWord.CompressedWord
        3
        ( case component of
            X -> 0
            Y -> 1
            Z -> 2
            W -> 3
        )
    )

putParts :: Double -> Double -> Double -> BitPut.BitPut
putParts a b c = putPart a <> putPart b <> putPart c

putPart :: Double -> BitPut.BitPut
putPart = CompressedWord.bitPut . compressPart

bitGet :: BitGet.BitGet Quaternion
bitGet = do
  component <- decodeComponent
  a <- decodePart
  b <- decodePart
  c <- decodePart
  pure $ toQuaternion component a b c

decodeComponent :: BitGet.BitGet Component
decodeComponent = do
  x_ <- CompressedWord.bitGet 3
  case CompressedWord.value x_ of
    0 -> pure X
    1 -> pure Y
    2 -> pure Z
    3 -> pure W
    y_ -> BitGet.throw $ InvalidComponent.InvalidComponent y_

decodePart :: BitGet.BitGet Double
decodePart = fmap decompressPart $ CompressedWord.bitGet maxCompressedValue