    -- |
    -- 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 :: forall a. (Enum a, Num a) => Infinite a
    pents = Infinite a -> Infinite a -> Infinite a
    forall a. Infinite a -> Infinite a -> Infinite a
    Inf.interleave
      ((a -> a -> a) -> a -> Infinite a -> Infinite a
    forall b a. (b -> a -> b) -> b -> Infinite a -> Infinite b
    Inf.scanl (\a
    acc a
    n -> a
    acc a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    3 a -> a -> a
    forall a. Num a => a -> a -> a
    * a
    n a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    1) a
    0 (a
    1...))
      ((a -> a -> a) -> a -> Infinite a -> Infinite a
    forall b a. (b -> a -> b) -> b -> Infinite a -> Infinite b
    Inf.scanl (\a
    acc a
    n -> a
    acc a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    3 a -> a -> a
    forall a. Num a => a -> a -> a
    * a
    n a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    2) a
    1 (a
    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 :: forall a. Num a => (Word -> a) -> Word -> a
    partitionF Word -> a
    _ Word
    0 = a
    1
    partitionF Word -> a
    f Word
    n
      = [a] -> a
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum
      ([a] -> a) -> [a] -> a
    forall a b. (a -> b) -> a -> b
    $ (a -> a -> a) -> [a] -> [a] -> [a]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith a -> a -> a
    forall a. Num a => a -> a -> a
    (*) ([a] -> [a]
    forall a. HasCallStack => [a] -> [a]
    cycle [a
    1, a
    1, -a
    1, -a
    1])
      ([a] -> [a]) -> [a] -> [a]
    forall a b. (a -> b) -> a -> b
    $ (Word -> a) -> [Word] -> [a]
    forall a b. (a -> b) -> [a] -> [b]
    map (Word -> a
    f (Word -> a) -> (Word -> Word) -> Word -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Word
    n -))
      ([Word] -> [a]) -> [Word] -> [a]
    forall a b. (a -> b) -> a -> b
    $ (Word -> Bool) -> Infinite Word -> [Word]
    forall a. (a -> Bool) -> Infinite a -> [a]
    Inf.takeWhile (Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Word
    n)
      (Infinite Word -> [Word]) -> Infinite Word -> [Word]
    forall a b. (a -> b) -> a -> b
    $ Infinite Word -> Infinite Word
    forall a. Infinite a -> Infinite a
    Inf.tail Infinite Word
    forall a. (Enum a, Num a) => Infinite a
    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 :: forall a. Num a => Infinite a
    partition = (Word -> a) -> Infinite a
    forall a. (Word -> a) -> Infinite a
    Inf.tabulate (Chimera Vector a -> Word -> a
    forall (v :: * -> *) a. Vector v a => Chimera v a -> Word -> a
    Ch.index Chimera Vector a
    ch)
      where
        ch :: Chimera Vector a
    ch = forall (v :: * -> *) a.
    Vector v a =>
    ((Word -> a) -> Word -> a) -> Chimera v a
    Ch.tabulateFix @Vector (Word -> a) -> Word -> a
    forall a. Num a => (Word -> a) -> Word -> a
    partitionF
    {-# SPECIALIZE partition :: Infinite Int     #-}
    {-# SPECIALIZE partition :: Infinite Word    #-}
    {-# SPECIALIZE partition :: Infinite Integer #-}
    {-# SPECIALIZE partition :: Infinite Natural #-}