{-# LANGUAGE TypeApplications #-}
module Main where

import Data.Foldable (toList)
import Hedgehog
import Test.Tasty
import Test.Tasty.Hedgehog
import Test.Tasty.HUnit

import qualified Numeric.AD as AD
import qualified Numeric.AD.Double as AD.Double

import Numeric.ADDual
import qualified Numeric.ADDual.Array.Internal as ADDA
import Numeric.ADDual.Examples


(~=) :: (Fractional a, Ord a) => a -> a -> Bool
x ~= y = abs (x - y) < 1e-5 || (let m = max (abs x) (abs y) in m > 1e-5 && abs (x - y) / m < 1e-5)

(~==) :: (Fractional a, Ord a, Show a) => a -> a -> PropertyT IO ()
x ~== y = diff x (~=) y

(~=!) :: (Foldable t, Fractional a, Ord a) => t a -> t a -> Bool
a ~=! b = length (toList a) == length (toList b) && and (zipWith (~=) (toList a) (toList b))

(~==!) :: (Foldable t, Fractional a, Ord a, Show (t a)) => t a -> t a -> PropertyT IO ()
a ~==! b = diff a (~=!) b


main :: IO ()
main = defaultMain $ testGroup "Tests"
  [testCase "product [1..5]" $
      gradient' @Double product [1..5] 1 @?= (120, [120, 60, 40, 30, 24])

  ,testProperty "neural 80" $ property $ do
      input <- forAll (genNeuralInput 80)
      let (res, grad) = gradient' fneural input 1
      res === fneural input
      grad ~==! AD.grad fneural input
      AD.grad fneural input === AD.Double.grad fneural input

  ,testProperty "neural 150" $ property $ do
      input <- forAll (genNeuralInput 150)
      let (res, grad) = gradient' fneural input 1
      res === fneural input
      grad ~==! AD.grad fneural input
      AD.grad fneural input === AD.Double.grad fneural input
    
  ,testProperty "primal neural == neural_A" $ property $ do
      input <- forAll (genNeuralInput_A 100)
      let resA = fneural_A input
      let res = fneural (cvtFNeuralAtoFNeural input)
      resA ~== res

  ,testProperty "neural_A 100" $ property $ do
      input <- forAll (genNeuralInput_A 100)
      let (resA, gradA) = ADDA.gradient' fneural_A input 1
      let (res, grad) = gradient' fneural (cvtFNeuralAtoFNeural input) 1
      resA ~== res
      cvtFNeuralAtoFNeural gradA ~==! grad
  ]