{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE TemplateHaskell                          #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Tests.EMD (
    emdTests
  ) where

import           Control.Monad
import           Data.Functor.Identity
import           Data.Proxy
import           GHC.TypeNats
import           Hedgehog
import           Numeric.EMD
import           Test.Tasty
import           Tests.Util
import qualified Hedgehog.Range        as Range

emdTests :: TestTree
emdTests = groupTree $$(discover)

prop_iemd_default :: Property
prop_iemd_default = iemdProp defaultEO

prop_orthog_default :: Property
prop_orthog_default = orthogProp defaultEO

edtEO :: EMDOpts Double
edtEO = defaultEO
    { eoSiftCondition = scEnergyDiff 0.01 0.01
                 `SCOr` SCTimes 100
    }

prop_iemd_edt :: Property
prop_iemd_edt = iemdProp edtEO

prop_orthog_edt :: Property
prop_orthog_edt = orthogProp edtEO

sCondEO :: EMDOpts Double
sCondEO = defaultEO
    { eoSiftCondition = SCSCond 10
                 `SCOr` SCTimes 100
    }

prop_iemd_sCond :: Property
prop_iemd_sCond = iemdProp sCondEO

prop_orthog_sCond :: Property
prop_orthog_sCond = orthogProp sCondEO


iemdProp :: EMDOpts Double -> Property
iemdProp eo = property $ withSize (Range.linear 1 8) $ \(_ :: Proxy n) -> do
    xs <- forAll $ generateData @n
    tripping (CE xs) (emd @_ @_ @(2^n-1) eo . getCE) (Identity . CE . iemd)

orthogProp :: EMDOpts Double -> Property
orthogProp eo = property $ withSize (Range.linear 8 10) $ \(_ :: Proxy n) -> do
    xs   <- forAll $ generateData @n
    let imfs = emdIMFs (emd @_ @_ @(2^n-1) eo xs)
        orthoMatrix =
          [ ((i, j), (x, y), dot x y / sqrt (dot x x * dot y y))
          | (i, x) <- zip indices imfs
          , (j, y) <- zip indices imfs
          , i < j
          ]
        badOrthos = filter (\(_,_,d) -> abs d > 0.5) orthoMatrix
        fracBad :: Double
        fracBad = fromIntegral (length badOrthos)
                / fromIntegral (length orthoMatrix)
    annotateShow orthoMatrix
    annotateShow fracBad
    when (length orthoMatrix < 6) discard
    assert $ fracBad <= 0.5
  where
    indices :: [Int]
    indices = [1..]