-- |
-- 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 TypeApplications #-}

module Math.NumberTheory.Recurrences.Pentagonal
  ( partition
  ) where

import qualified Data.Chimera as Ch
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) => [a]
pents :: [a]
pents = [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave ((a -> a -> a) -> a -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
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 -> [a] -> [a]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
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..])
  where
    interleave :: [a] -> [a] -> [a]
    interleave :: [a] -> [a] -> [a]
interleave (a
n : [a]
ns) (a
m : [a]
ms) = a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: a
m a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
interleave [a]
ns [a]
ms
    interleave [a]
_ [a]
_ = []

-- | @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 :: (Word -> a) -> Word -> a
partitionF Word -> a
_ Word
0 = a
1
partitionF Word -> a
f Word
n
  = [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. [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 -> Word -> Word
forall a. Num a => a -> a -> a
-))
  ([Word] -> [a]) -> [Word] -> [a]
forall a b. (a -> b) -> a -> b
$ (Word -> Bool) -> [Word] -> [Word]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
<= Word
n)
  ([Word] -> [Word]) -> [Word] -> [Word]
forall a b. (a -> b) -> a -> b
$ [Word] -> [Word]
forall a. [a] -> [a]
tail [Word]
forall a. (Enum a, Num a) => [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 => [a]
partition :: [a]
partition = Chimera Vector a -> [a]
forall (v :: * -> *) a. Vector v a => Chimera v a -> [a]
Ch.toList (Chimera Vector a -> [a]) -> Chimera Vector a -> [a]
forall a b. (a -> b) -> a -> b
$ ((Word -> a) -> Word -> a) -> Chimera Vector a
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 :: [Int]     #-}
{-# SPECIALIZE partition :: [Word]    #-}
{-# SPECIALIZE partition :: [Integer] #-}
{-# SPECIALIZE partition :: [Natural] #-}