content: -- |
  -- Module:      Math.NumberTheory.Recurrences.Pentagonal
  -- Copyright:   (c) 2018 Alexandre Rodrigues Baldé
  -- Licence:     MIT
  -- Maintainer:  Alexandre Rodrigues Baldé <alexandrer_b@outlook.com>
  --
  -- Values of <https://en.wikipedia.org/wiki/Partition_(number_theory)#Partition_function partition function>.
  --
  
  {-# LANGUAGE PostfixOperators #-}
  {-# LANGUAGE TypeApplications #-}
  
  module Math.NumberTheory.Recurrences.Pentagonal
    ( partition
    ) where
  
  import qualified Data.Chimera as Ch
  import Data.List.Infinite (Infinite(..), (...))
  import qualified Data.List.Infinite as Inf
  import Data.Vector (Vector)
  import Numeric.Natural (Natural)
  
  -- | Infinite list of generalized pentagonal numbers.
  -- Example:
  --
  -- >>> take 10 pents
  -- [0,1,2,5,7,12,15,22,26,35]
  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...))
  
  -- | @p(n) = p(n-1) + p(n-2) - p(n-5) - p(n-7) + p(n-11) + ...@, where @p(0) = 1@
  -- and @p(k) = 0@ for a negative integer @k@. Uses a @Chimera@ from the
  -- @chimera@ package to memoize previous results.
  partitionF :: Num a => (Word -> a) -> Word -> a
  partitionF _ 0 = 1
  partitionF f n
    = sum
    $ zipWith (*) (cycle [1, 1, -1, -1])
    $ map (f . (n -))
    $ Inf.takeWhile (<= n)
    $ Inf.tail pents
  
  -- | Infinite zero-based table of <https://oeis.org/A000041 partition numbers>.
  --
  -- >>> take 10 partition
  -- [1,1,2,3,5,7,11,15,22,30]
  --
  -- >>> :set -XDataKinds
  -- >>> import Data.Mod
  -- >>> partition !! 1000 :: Mod 1000
  -- (991 `modulo` 1000)
  partition :: Num a => Infinite a
  partition = Inf.tabulate (Ch.index ch)
    where
      ch = Ch.tabulateFix @Vector partitionF
  {-# SPECIALIZE partition :: Infinite Int     #-}
  {-# SPECIALIZE partition :: Infinite Word    #-}
  {-# SPECIALIZE partition :: Infinite Integer #-}
  {-# SPECIALIZE partition :: Infinite Natural #-}
  