{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}

module Spec.Fusion (tests) where

import Control.DeepSeq (NFData, force)
import Control.Concurrent (forkIO)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Exception (SomeException, evaluate, throwIO, try)
import Data.Int (Int64)
import Data.List qualified as List
import Data.Word (Word32, Word64)
import GHC.Float (castDoubleToWord64, castFloatToWord32)
import Hedgehog
import System.Mem (getAllocationCounter)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testPropertyNamed)

import Geomancy.DVec4 (DVec4, dvec4, withDVec4)
import Geomancy.DVec4 qualified as DVec4
import Geomancy.Swizzle ((.@))
import Geomancy.Vec4 (Vec4, vec4, withVec4)
import Geomancy.Vec4 qualified as Vec4

import Support.Gen (genDVec4, genDouble, genFloat, genVec4)
import Support.Property (pattern PROP_TESTS)

tests :: TestTree
tests = testGroup "Fusion"
  [ testPropertyNamed "vec4_bits" "prop_vec4_fusion_bits" prop_vec4_fusion_bits
  , testPropertyNamed "dvec4_bits" "prop_dvec4_fusion_bits" prop_dvec4_fusion_bits
  , testPropertyNamed "vec4_alloc" "prop_vec4_fusion_alloc" prop_vec4_fusion_alloc
  , testPropertyNamed "dvec4_alloc" "prop_dvec4_fusion_alloc" prop_dvec4_fusion_alloc
  ]

{-# NOINLINE unfused #-}
unfused :: a -> a
unfused x = x

bits4 :: Vec4 -> (Word32, Word32, Word32, Word32)
bits4 v = withVec4 v \x y z w ->
  (castFloatToWord32 x, castFloatToWord32 y, castFloatToWord32 z, castFloatToWord32 w)

dbits4 :: DVec4 -> (Word64, Word64, Word64, Word64)
dbits4 v = withDVec4 v \x y z w ->
  (castDoubleToWord64 x, castDoubleToWord64 y, castDoubleToWord64 z, castDoubleToWord64 w)

prop_vec4_fusion_bits :: Property
prop_vec4_fusion_bits = withTests PROP_TESTS $ property do
  t <- forAll genFloat
  a <- forAll genVec4
  b <- forAll genVec4
  c <- forAll genVec4
  d <- forAll genVec4
  bits4 (Vec4.lerp t a b) ===
    bits4 (unfused (unfused (a Vec4.^* t) + unfused (b Vec4.^* (1 - t))))
  bits4 ((a + b) * c - d) ===
    bits4 (unfused (unfused (unfused (a + b) * c) - d))
  castFloatToWord32 (Vec4.dot (a + b) c) ===
    castFloatToWord32 (Vec4.dot (unfused (a + b)) c)
  bits4 (Vec4.normalize (a + b)) ===
    bits4 (unfused (Vec4.normalize (unfused (a + b))))
  bits4 (sqrt (a * b)) ===
    bits4 (unfused (sqrt (unfused (a * b))))
  bits4 ((a + b) .@ #wzyx) ===
    bits4 (unfused (unfused (a + b) .@ #wzyx))
  castFloatToWord32 ((a + b).wzyx.wzyx.y + c.xyzz.w) ===
    castFloatToWord32 ((unfused (unfused (unfused (a + b)).wzyx).wzyx).y + (unfused c.xyzz).w)
  castFloatToWord32 ((a + b).wzyx.zyx.yx.y) ===
    castFloatToWord32 ((unfused (unfused (unfused (unfused (a + b)).wzyx).zyx).yx).y)
  bits4 ((a + b).xyz.xyz1 + c.xy.xy01) ===
    bits4 (unfused (unfused (unfused (a + b)).xyz).xyz1 + unfused (unfused c.xy).xy01)
  bits4 (a + b).yzw.xyz0.wxyz ===
    bits4 (withVec4 (unfused (a + b)) \_ y z w -> vec4 0 y z w)

prop_dvec4_fusion_bits :: Property
prop_dvec4_fusion_bits = withTests PROP_TESTS $ property do
  t <- forAll genDouble
  a <- forAll genDVec4
  b <- forAll genDVec4
  c <- forAll genDVec4
  d <- forAll genDVec4
  dbits4 (DVec4.lerp t a b) ===
    dbits4 (unfused (unfused (a DVec4.^* t) + unfused (b DVec4.^* (1 - t))))
  dbits4 ((a + b) * c - d) ===
    dbits4 (unfused (unfused (unfused (a + b) * c) - d))
  castDoubleToWord64 (DVec4.dot (a + b) c) ===
    castDoubleToWord64 (DVec4.dot (unfused (a + b)) c)
  dbits4 (DVec4.normalize (a + b)) ===
    dbits4 (unfused (DVec4.normalize (unfused (a + b))))
  dbits4 (sqrt (a * b)) ===
    dbits4 (unfused (sqrt (unfused (a * b))))
  castDoubleToWord64 ((a + b).wzyx.wzyx.y + c.xyzz.w) ===
    castDoubleToWord64 ((unfused (unfused (unfused (a + b)).wzyx).wzyx).y + (unfused c.xyzz).w)
  castDoubleToWord64 ((a + b).wzyx.zyx.yx.y) ===
    castDoubleToWord64 ((unfused (unfused (unfused (unfused (a + b)).wzyx).zyx).yx).y)
  dbits4 ((a + b).xyz.xyz1 + c.xy.xy01) ===
    dbits4 (unfused (unfused (unfused (a + b)).xyz).xyz1 + unfused (unfused c.xy).xy01)
  dbits4 (a + b).yzw.xyz0.wxyz ===
    dbits4 (withDVec4 (unfused (a + b)) \_ y z w -> dvec4 0 y z w)

samples :: Int
samples = 1000

scalarBudget :: Int64
scalarBudget = 1024

vectorBudget :: Int64 -> Int64
vectorBudget bytesPerVector = bytesPerVector * fromIntegral samples + scalarBudget

allocated :: (a -> b) -> a -> IO Int64
allocated f x = do
  result <- newEmptyMVar
  _ <- forkIO $ try (allocatedOnFreshStack f x) >>= putMVar result
  takeMVar result >>= either rethrow pure
  where
    rethrow :: SomeException -> IO Int64
    rethrow = throwIO

allocatedOnFreshStack :: (a -> b) -> a -> IO Int64
allocatedOnFreshStack f x = do
  before <- getAllocationCounter
  _ <- evaluate (f x)
  after <- getAllocationCounter
  pure (before - after)

inputs :: NFData v => (Float -> v) -> IO [v]
inputs mk = evaluate $ force
  [ mk (fromIntegral i / fromIntegral samples)
  | i <- [1 .. samples]
  ]

overBudget :: Int64 -> [(String, Int64)] -> [(String, Int64)]
overBudget budget = filter \(_name, bytes) -> bytes > budget

{-# NOINLINE plainFloatSum #-}
plainFloatSum :: [Float] -> Float
plainFloatSum = List.foldl' (+) 0

builtOptimized :: IO Bool
builtOptimized = do
  floats <- inputs id
  bytes <- allocated plainFloatSum floats
  pure $ bytes <= scalarBudget

whenBuiltOptimized :: PropertyT IO () -> PropertyT IO ()
whenBuiltOptimized checks = do
  optimized <- evalIO builtOptimized
  if optimized then
    checks
  else
    annotate "Skipped: allocation budgets only hold in optimized builds"

prop_vec4_fusion_alloc :: Property
prop_vec4_fusion_alloc = withTests 1 $ property $ whenBuiltOptimized do
  vs <- evalIO $ inputs \f -> vec4 f (f * 0.5) (f + 1) (2 - f)
  let
    b = vec4 1 2 3 4
    c = vec4 0.5 0.25 0.125 1
    d = vec4 3 2 1 0
  scalars <- evalIO $ traverse (traverse ($ vs))
    [ ("dot of sum", allocated (vec4DotSum b c))
    , ("normalize", allocated vec4NormalizeDot)
    , ("sqrt", allocated vec4SqrtDot)
    , ("swizzle read", allocated vec4SwizzleRead)
    , ("field swizzles", allocated vec4FieldSwizzles)
    , ("field extensions", allocated vec4FieldExtensions)
    ]
  vectors <- evalIO $ traverse (traverse ($ vs))
    [ ("lerp", allocated vec4Lerp)
    , ("chain", allocated (vec4Chain b c d))
    ]
  overBudget scalarBudget scalars ++ overBudget (vectorBudget 32) vectors === []

prop_dvec4_fusion_alloc :: Property
prop_dvec4_fusion_alloc = withTests 1 $ property $ whenBuiltOptimized do
  vs <- evalIO $ inputs \f ->
    let f' = realToFrac f in dvec4 f' (f' * 0.5) (f' + 1) (2 - f')
  let
    b = dvec4 1 2 3 4
    c = dvec4 0.5 0.25 0.125 1
    d = dvec4 3 2 1 0
  scalars <- evalIO $ traverse (traverse ($ vs))
    [ ("dot of sum", allocated (dvec4DotSum b c))
    , ("normalize", allocated dvec4NormalizeDot)
    , ("sqrt", allocated dvec4SqrtDot)
    , ("field swizzles", allocated dvec4FieldSwizzles)
    , ("field extensions", allocated dvec4FieldExtensions)
    ]
  vectors <- evalIO $ traverse (traverse ($ vs))
    [ ("lerp", allocated dvec4Lerp)
    , ("chain", allocated (dvec4Chain b c d))
    ]
  overBudget scalarBudget scalars ++ overBudget (vectorBudget 48) vectors === []

{-# NOINLINE vec4DotSum #-}
vec4DotSum :: Vec4 -> Vec4 -> [Vec4] -> Float
vec4DotSum b c = List.foldl' (\s v -> s + Vec4.dot (v + b) c) 0

{-# NOINLINE vec4NormalizeDot #-}
vec4NormalizeDot :: [Vec4] -> Float
vec4NormalizeDot = List.foldl' (\s v -> s + Vec4.dot (Vec4.normalize v) v) 0

{-# NOINLINE vec4SqrtDot #-}
vec4SqrtDot :: [Vec4] -> Float
vec4SqrtDot = List.foldl' (\s v -> s + Vec4.dot (sqrt v) v) 0

{-# NOINLINE vec4SwizzleRead #-}
vec4SwizzleRead :: [Vec4] -> Float
vec4SwizzleRead = List.foldl' (\s v -> s + withVec4 (v .@ #wzyx) \x _ _ _ -> x) 0

{-# NOINLINE vec4FieldSwizzles #-}
vec4FieldSwizzles :: [Vec4] -> Float
vec4FieldSwizzles = List.foldl' (\s v -> s + v.wzyx.wzyx.y + v.xyzz.w + v.wzyx.zyx.yx.y) 0

{-# NOINLINE vec4FieldExtensions #-}
vec4FieldExtensions :: [Vec4] -> Float
vec4FieldExtensions = List.foldl' (\s v -> s + Vec4.dot v.xyz.xyz1 v + Vec4.dot v.wz.xy01 v) 0

{-# NOINLINE vec4Lerp #-}
vec4Lerp :: [Vec4] -> Vec4
vec4Lerp = List.foldl' (\acc v -> Vec4.lerp 0.3 acc v) 0

{-# NOINLINE vec4Chain #-}
vec4Chain :: Vec4 -> Vec4 -> Vec4 -> [Vec4] -> Vec4
vec4Chain b c d = List.foldl' (\acc v -> acc + ((v + b) * c - d)) 0

{-# NOINLINE dvec4DotSum #-}
dvec4DotSum :: DVec4 -> DVec4 -> [DVec4] -> Double
dvec4DotSum b c = List.foldl' (\s v -> s + DVec4.dot (v + b) c) 0

{-# NOINLINE dvec4NormalizeDot #-}
dvec4NormalizeDot :: [DVec4] -> Double
dvec4NormalizeDot = List.foldl' (\s v -> s + DVec4.dot (DVec4.normalize v) v) 0

{-# NOINLINE dvec4SqrtDot #-}
dvec4SqrtDot :: [DVec4] -> Double
dvec4SqrtDot = List.foldl' (\s v -> s + DVec4.dot (sqrt v) v) 0

{-# NOINLINE dvec4FieldSwizzles #-}
dvec4FieldSwizzles :: [DVec4] -> Double
dvec4FieldSwizzles = List.foldl' (\s v -> s + v.wzyx.wzyx.y + v.xyzz.w + v.wzyx.zyx.yx.y) 0

{-# NOINLINE dvec4FieldExtensions #-}
dvec4FieldExtensions :: [DVec4] -> Double
dvec4FieldExtensions = List.foldl' (\s v -> s + DVec4.dot v.xyz.xyz1 v + DVec4.dot v.wz.xy01 v) 0

{-# NOINLINE dvec4Lerp #-}
dvec4Lerp :: [DVec4] -> DVec4
dvec4Lerp = List.foldl' (\acc v -> DVec4.lerp 0.3 acc v) 0

{-# NOINLINE dvec4Chain #-}
dvec4Chain :: DVec4 -> DVec4 -> DVec4 -> [DVec4] -> DVec4
dvec4Chain b c d = List.foldl' (\acc v -> acc + ((v + b) * c - d)) 0