{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilies                             #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}


module Tests.Util (
    groupTree
  , CloseEnough(..)
  , generateData
  , withSize
  , dot
  ) where

import           Data.Complex
import           Data.Proxy
import           Data.Type.Equality
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Hedgehog
import           Hedgehog.Internal.Property
import           Numeric.Natural
import           Statistics.Transform
import           Test.Tasty
import           Test.Tasty.Hedgehog
import qualified Data.Vector.Sized          as V
import qualified Hedgehog.Gen               as Gen
import qualified Hedgehog.Range             as Range

groupTree :: Group -> TestTree
groupTree Group{..} = testGroup (unGroupName groupName)
                                (map (uncurry go) groupProperties)
  where
    go :: PropertyName -> Property -> TestTree
    go n = testProperty (mkName (unPropertyName n))
    mkName = map deUnderscore . drop (length @[] @Char "prop_")
    deUnderscore '_' = ' '
    deUnderscore c   = c

newtype CloseEnough n = CE { getCE :: V.Vector n Double }
  deriving Show

instance KnownNat n => Eq (CloseEnough n) where
    CE x == CE y = ((d `dot` d) / sqrt ((x `dot` x) * (y `dot` y))) < 0.0001
      where
        d = V.zipWith (-) x y

dot :: Num a => V.Vector n a -> V.Vector n a -> a
dot a b = sum $ V.zipWith (*) a b

withSize
    :: Monad m
    => Range Natural
    -> (forall n. (KnownNat n, 1 <= 2^n) => Proxy n -> PropertyT m a)
    -> PropertyT m a
withSize r f = do
    n <- forAll $ Gen.integral r
    case someNatVal n of
      SomeNat (p :: Proxy n) -> do
        LE Refl <- pure $ Proxy @1 %<=? Proxy @(2^n)
        f p

generateData
    :: KnownNat n
    => Gen (V.Vector (2^n) Double)
generateData = fmap (fmap realPart . ifftSized) . V.generateM $ \i ->
    let i' = recip . (+ 1) . fromIntegral $ i
    in  mkPolar <$> Gen.double (Range.exponentialFloat (i' / 10) i')
                <*> Gen.double (Range.constant (-pi) pi)

ifftSized
    :: V.Vector (2^n) (Complex Double)
    -> V.Vector (2^n) (Complex Double)
ifftSized = V.withVectorUnsafe ifft