-- |
-- 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, Typeable v) =>
((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 #-}