content: -- |
  -- Module:      Math.NumberTheory.Recurrences.PentagonalTests
  -- Copyright:   (c) 2018 Alexandre Rodrigues Baldé
  -- Licence:     MIT
  -- Maintainer:  Alexandre Rodrigues Baldé <alexandrer_b@outlook.com>
  --
  -- Tests for Math.NumberTheory.Recurrences.Pentagonal
  --
  
  {-# LANGUAGE PostfixOperators    #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE ViewPatterns        #-}
  
  {-# OPTIONS_GHC -fno-warn-type-defaults #-}
  
  module Math.NumberTheory.Recurrences.PentagonalTests
    ( testSuite
    ) where
  
  import Data.List.Infinite (Infinite(..), (...))
  import qualified Data.List.Infinite as Inf
  import Data.Proxy                    (Proxy (..))
  import GHC.Natural                   (Natural)
  import GHC.TypeNats                  (SomeNat (..), someNatVal)
  
  import Math.NumberTheory.Moduli      (Mod, getVal)
  import Math.NumberTheory.Recurrences (partition)
  import Math.NumberTheory.TestUtils
  
  import Test.Tasty
  import Test.Tasty.HUnit
  
  -- | Helper to avoid writing @partition !!@ too many times.
  partition' :: Num a => Int -> a
  partition' = (partition Inf.!!) . fromIntegral
  
  -- | Check that the @k@-th generalized pentagonal number is
  -- @div (3 * k² - k) 2@, where @k ∈ {0, 1, -1, 2, -2, 3, -3, 4, ...}@.
  -- Notice that @-1@ is the @2 * abs (-1) == 2@-nd index in the zero-based list,
  -- while @2@ is the @2 * 2 - 1 == 3@-rd, and so on.
  pentagonalNumbersProperty1 :: AnySign Int -> Bool
  pentagonalNumbersProperty1 (AnySign n)
      | n == 0    = Inf.head pents           == 0
      | n > 0     = pents Inf.!! (2 * fromIntegral n - 1) == pent n
      | otherwise = pents Inf.!! (2 * fromIntegral (- n)) == pent n
    where
      pent m = div (3 * (m * m) - m) 2
  
  -- | Check that the first 20 elements of @partition@ are correct per
  -- https://oeis.org/A000041.
  partitionSpecialCase20 :: Assertion
  partitionSpecialCase20 = assertEqual "partition"
      (Inf.take 20 partition)
      [1, 1, 2, 3, 5, 7, 11, 15, 22, 30, 42, 56, 77, 101, 135, 176, 231, 297, 385, 490]
  
  -- | Copied from @Math.NumberTheory.Recurrences.Pentagonal@ to test the
  -- reference implementation of @partition@.
  pentagonalSigns :: Num a => [a] -> [a]
  pentagonalSigns = zipWith (*) (cycle [1, 1, -1, -1])
  
  -- | Copied from @Math.NumberTheory.Recurrences.Pentagonal@ to test the
  -- reference implementation of @partition@.
  pents :: (Enum a, Num a) => Infinite a
  pents = Inf.interleave
    (Inf.scanl (\acc n -> acc + 3 * n - 1) 0 (1...))
    (Inf.scanl (\acc n -> acc + 3 * n - 2) 1 (2...))
  
  -- | Check that @p(n) = p(n-1) + p(n-2) - p(n-5) - p(n-7) + p(n-11) + ...@,
  -- where @p(x) = 0@ for any negative integer and @p(0) = 1@.
  partitionProperty1 :: Positive Int -> Bool
  partitionProperty1 (Positive n) =
      partition' n == (sum .
                       pentagonalSigns .
                       map (\m -> partition' (n - m)) .
                       Inf.takeWhile (\m -> n - m >= 0) .
                       Inf.tail $ pents)
  
  -- | Check that
  -- @partition :: [Math.NumberTheory.Moduli.Mod n] == map (`mod` n) partition@.
  partitionProperty2 :: NonNegative Integer -> Positive Natural -> Bool
  partitionProperty2 (NonNegative m)
                     n@(someNatVal . getPositive -> (SomeNat (Proxy :: Proxy n))) =
      (Inf.take m' . Inf.map getVal $ (partition :: Infinite (Mod n))) ==
      map helper (Inf.take m' partition :: [Integer])
    where
      m' = fromIntegral m
      n' = fromIntegral n
      helper x = x `mod` n'
  
  testSuite :: TestTree
  testSuite = testGroup "Pentagonal"
    [ testGroup "partition"
      [ testSmallAndQuick "matches definition"  partitionProperty1
      , testSmallAndQuick "mod n" partitionProperty2
      , testCase          "first 20 elements of partition are correct"
                          partitionSpecialCase20
      ]
    , testGroup "Generalized pentagonal numbers"
      [ testSmallAndQuick "matches definition" pentagonalNumbersProperty1
      ]
    ]
  