{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
-- {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Tests.C where

import Control.Monad
import Data.Array.RankedS qualified as OR
import Data.Foldable (toList)
import Data.Type.Equality
import Foreign
import GHC.TypeLits

import Data.Array.Mixed.XArray qualified as X
import Data.Array.Mixed.Lemmas
import Data.Array.Nested
import Data.Array.Nested.Internal.Mixed
import Data.Array.Nested.Internal.Shape

import Hedgehog
import Hedgehog.Internal.Property (forAllT)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Test.Tasty
import Test.Tasty.Hedgehog

-- import Debug.Trace

import Gen
import Util


tests :: TestTree
tests = testGroup "C"
  [testGroup "sum"
    [testProperty "random nonempty" $ property $ genRank $ \outrank@(SNat @n) -> do
      -- Test nonempty _results_. The first dimension of the input is allowed to be 0, because then OR.rerank doesn't fail yet.
      let inrank = SNat @(n + 1)
      sh <- forAll $ genShR inrank
      -- traceM ("sh: " ++ show sh ++ " -> " ++ show (product sh))
      guard (all (> 0) (toList (shrTail sh)))  -- only constrain the tail
      arr <- forAllT $ OR.fromVector @Double @(n + 1) (toList sh) <$>
               genStorables (Range.singleton (product sh))
                            (\w -> fromIntegral w / fromIntegral (maxBound :: Word64))
      let rarr = rfromOrthotope inrank arr
      -- annotateShow rarr
      Refl <- return $ lemRankReplicate outrank
      let Ranked (M_Double (M_Primitive _ (X.XArray lhs))) = rsumOuter1 rarr
      let rhs = orSumOuter1 outrank arr
      -- annotateShow lhs
      -- annotateShow rhs
      lhs === rhs

    ,testProperty "random empty" $ property $ genRank $ \outrankm1@(SNat @nm1) -> do
      -- We only need to test shapes where the _result_ is empty; the rest is handled by 'random nonempty' above.
      outrank :: SNat n <- return $ SNat @(nm1 + 1)
      let inrank = SNat @(n + 1)
      sh <- forAll $ do
        shtt <- genShR outrankm1  -- nm1
        sht <- shuffleShR (0 :$: shtt)  -- n
        n <- Gen.int (Range.linear 0 20)
        return (n :$: sht)  -- n + 1
      guard (any (== 0) (toList (shrTail sh)))
      -- traceM ("sh: " ++ show sh ++ " -> " ++ show (product sh))
      let arr = OR.fromList @Double @(n + 1) (toList sh) []
      let rarr = rfromOrthotope inrank arr
      Refl <- return $ lemRankReplicate outrank
      let Ranked (M_Double (M_Primitive _ (X.XArray lhs))) = rsumOuter1 rarr
      OR.toList lhs === []
    ]
  ]