-- | Counting partitions of integers.

{-# LANGUAGE CPP, BangPatterns, ScopedTypeVariables #-}
module Math.Combinat.Partitions.Integer.Count where

--------------------------------------------------------------------------------

import Data.List
import Control.Monad ( liftM , replicateM )

-- import Data.Map (Map)
-- import qualified Data.Map as Map

import Math.Combinat.Numbers ( factorial , binomial , multinomial )
import Math.Combinat.Numbers.Integers -- Primes
import Math.Combinat.Helper

import Data.Array
import System.Random

--------------------------------------------------------------------------------
-- * Infinite tables of integers

-- | A data structure which is essentially an infinite list of @Integer@-s,
-- but fast lookup (for reasonable small inputs)
newtype TableOfIntegers = TableOfIntegers [Array Int Integer]

lookupInteger :: TableOfIntegers -> Int -> Integer
lookupInteger :: TableOfIntegers -> Int -> Integer
lookupInteger (TableOfIntegers [Array Int Integer]
table) !Int
n 
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0  = ([Array Int Integer]
table [Array Int Integer] -> Int -> Array Int Integer
forall a. [a] -> Int -> a
!! Int
k) Array Int Integer -> Int -> Integer
forall i e. Ix i => Array i e -> i -> e
! Int
r
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
0  = Integer
0
  where
    (Int
k,Int
r) = Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
divMod Int
n Int
1024

makeTableOfIntegers
  :: ((Int -> Integer) -> (Int -> Integer))
  -> TableOfIntegers
makeTableOfIntegers :: ((Int -> Integer) -> Int -> Integer) -> TableOfIntegers
makeTableOfIntegers (Int -> Integer) -> Int -> Integer
user = TableOfIntegers
table where
  calc :: Int -> Integer
calc  = (Int -> Integer) -> Int -> Integer
user Int -> Integer
lkp
  lkp :: Int -> Integer
lkp   = TableOfIntegers -> Int -> Integer
lookupInteger TableOfIntegers
table
  table :: TableOfIntegers
table = [Array Int Integer] -> TableOfIntegers
TableOfIntegers
    [ (Int, Int) -> [Integer] -> Array Int Integer
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0,Int
1023) ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
calc [Int
a..Int
b]) 
    | Int
k<-[Int
0..] 
    , let a :: Int
a = Int
1024Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
k 
    , let b :: Int
b = Int
1024Int -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 
    ]

--------------------------------------------------------------------------------
-- * Counting partitions

-- | Number of partitions of @n@ (looking up a table built using Euler's algorithm)
countPartitions :: Int -> Integer
countPartitions :: Int -> Integer
countPartitions = TableOfIntegers -> Int -> Integer
lookupInteger TableOfIntegers
partitionCountTable 

-- | This uses the power series expansion of the infinite product. It is slower than the above.
countPartitionsInfiniteProduct :: Int -> Integer
countPartitionsInfiniteProduct :: Int -> Integer
countPartitionsInfiniteProduct Int
k = [Integer]
partitionCountListInfiniteProduct [Integer] -> Int -> Integer
forall a. [a] -> Int -> a
!! Int
k

-- | This uses 'countPartitions'', and is (very) slow
countPartitionsNaive :: Int -> Integer
countPartitionsNaive :: Int -> Integer
countPartitionsNaive Int
d = (Int, Int) -> Int -> Integer
countPartitions' (Int
d,Int
d) Int
d

--------------------------------------------------------------------------------

-- | This uses Euler's algorithm to compute p(n)
--
-- See eg.:
-- NEIL CALKIN, JIMENA DAVIS, KEVIN JAMES, ELIZABETH PEREZ, AND CHARLES SWANNACK
-- COMPUTING THE INTEGER PARTITION FUNCTION
-- <http://www.math.clemson.edu/~kevja/PAPERS/ComputingPartitions-MathComp.pdf>
--
partitionCountTable :: TableOfIntegers
partitionCountTable :: TableOfIntegers
partitionCountTable = TableOfIntegers
table where

  table :: TableOfIntegers
table = ((Int -> Integer) -> Int -> Integer) -> TableOfIntegers
makeTableOfIntegers (Int -> Integer) -> Int -> Integer
forall p. Num p => (Int -> p) -> Int -> p
fun

  fun :: (Int -> p) -> Int -> p
fun Int -> p
lkp !Int
n 
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
1 = (p -> p -> p) -> p -> [p] -> p
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' p -> p -> p
forall a. Num a => a -> a -> a
(+) p
0 
             [ (if Int -> Bool
forall a. Integral a => a -> Bool
even Int
k then p -> p
forall a. Num a => a -> a
negate else p -> p
forall a. a -> a
id) 
                 ( Int -> p
lkp (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int
forall a. Integral a => a -> a -> a
div (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
3Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)) Int
2)
                 p -> p -> p
forall a. Num a => a -> a -> a
+ Int -> p
lkp (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int
forall a. Integral a => a -> a -> a
div (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
*(Int
3Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) Int
2)
                 )
             | Int
k <- [Int
1..Int -> Int
limit Int
n]
             ]
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
0 = p
0
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = p
1
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = p
1

  limit :: Int -> Int
  limit :: Int -> Int
limit !Int
n = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Integer -> Integer
ceilingSquareRoot (Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
nnInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
nnInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1) Integer
3) where
    nn :: Integer
nn = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n :: Integer

-- | An infinite list containing all @p(n)@, starting from @p(0)@.
partitionCountList :: [Integer]
partitionCountList :: [Integer]
partitionCountList = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
countPartitions [Int
0..]

--------------------------------------------------------------------------------

-- | Infinite list of number of partitions of @0,1,2,...@
--
-- This uses the infinite product formula the generating function of partitions, 
-- recursively expanding it; it is reasonably fast for small numbers.
--
-- > partitionCountListInfiniteProduct == map countPartitions [0..]
--
partitionCountListInfiniteProduct :: [Integer]
partitionCountListInfiniteProduct :: [Integer]
partitionCountListInfiniteProduct = [Integer]
final where

  final :: [Integer]
final = Int -> [Integer] -> [Integer]
go Int
1 (Integer
1Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
:Integer -> [Integer]
forall a. a -> [a]
repeat Integer
0) 

  go :: Int -> [Integer] -> [Integer]
go !Int
k (Integer
x:[Integer]
xs) = Integer
x Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Int -> [Integer] -> [Integer]
go (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Integer]
ys where
    ys :: [Integer]
ys = (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(+) [Integer]
xs (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
k [Integer]
final [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [Integer]
ys)
    -- explanation:
    --   xs == drop k $ f (k-1)
    --   ys == drop k $ f (k  )  

{-

Full explanation of 'partitionCountList':
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

let f k = productPSeries $ map (:[]) [1..k]

f 0 = [1,0,0,0,0,0,0,0...]
f 1 = [1,1,1,1,1,1,1,1...]
f 2 = [1,1,2,2,3,3,4,4...]
f 3 = [1,1,2,3,4,5,7,8...]

observe: 

* take (k+1) (f k) == take (k+1) partitionCountList
* f (k+1) == zipWith (+) (f k) (replicate (k+1) 0 ++ f (k+1))

now apply (drop (k+1)) to the second one : 

* drop (k+1) (f (k+1)) == zipWith (+) (drop (k+1) $ f k) (f (k+1))
* f (k+1) = take (k+1) final ++ drop (k+1) (f (k+1))

-}

--------------------------------------------------------------------------------

-- | Naive infinite list of number of partitions of @0,1,2,...@
--
-- > partitionCountListNaive == map countPartitionsNaive [0..]
--
-- This is very slow.
--
partitionCountListNaive :: [Integer]
partitionCountListNaive :: [Integer]
partitionCountListNaive = (Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
countPartitionsNaive [Int
0..]

--------------------------------------------------------------------------------
-- * Counting all partitions

countAllPartitions :: Int -> Integer
countAllPartitions :: Int -> Integer
countAllPartitions Int
d = [Integer] -> Integer
forall a. Num a => [a] -> a
sum' [ Int -> Integer
countPartitions Int
i | Int
i <- [Int
0..Int
d] ]

-- | Count all partitions fitting into a rectangle.
-- # = \\binom { h+w } { h }
countAllPartitions' :: (Int,Int) -> Integer
countAllPartitions' :: (Int, Int) -> Integer
countAllPartitions' (Int
h,Int
w) = 
  Int -> Int -> Integer
forall a. Integral a => a -> a -> Integer
binomial (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
w) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
h Int
w)
  --sum [ countPartitions' (h,w) i | i <- [0..d] ] where d = h*w

--------------------------------------------------------------------------------
-- * Counting fitting into a rectangle

-- | Number of of d, fitting into a given rectangle. Naive recursive algorithm.
countPartitions' :: (Int,Int) -> Int -> Integer
countPartitions' :: (Int, Int) -> Int -> Integer
countPartitions' (Int, Int)
_ Int
0 = Integer
1
countPartitions' (Int
0,Int
_) Int
d = if Int
dInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0 then Integer
1 else Integer
0
countPartitions' (Int
_,Int
0) Int
d = if Int
dInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0 then Integer
1 else Integer
0
countPartitions' (Int
h,Int
w) Int
d = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
  [ (Int, Int) -> Int -> Integer
countPartitions' (Int
i,Int
wInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i) | Int
i <- [Int
1..Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
d Int
h] ] 

--------------------------------------------------------------------------------
-- * Partitions with given number of parts

-- | Count partitions of @n@ into @k@ parts.
--
-- Naive recursive algorithm.
--
countPartitionsWithKParts 
  :: Int    -- ^ @k@ = number of parts
  -> Int    -- ^ @n@ = the integer we partition
  -> Integer
countPartitionsWithKParts :: Int -> Int -> Integer
countPartitionsWithKParts Int
k Int
n = Int -> Int -> Int -> Integer
forall t p. (Ord t, Num t, Num p, Enum t) => t -> t -> t -> p
go Int
n Int
k Int
n where
  go :: t -> t -> t -> p
go !t
h !t
k !t
n 
    | t
k t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<  t
0     = p
0
    | t
k t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
0     = if t
ht -> t -> Bool
forall a. Ord a => a -> a -> Bool
>=t
0 Bool -> Bool -> Bool
&& t
nt -> t -> Bool
forall a. Eq a => a -> a -> Bool
==t
0 then p
1 else p
0
    | t
k t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
1     = if t
ht -> t -> Bool
forall a. Ord a => a -> a -> Bool
>=t
n Bool -> Bool -> Bool
&& t
nt -> t -> Bool
forall a. Ord a => a -> a -> Bool
>=t
1 then p
1 else p
0
    | Bool
otherwise  = [p] -> p
forall a. Num a => [a] -> a
sum' [ t -> t -> t -> p
go t
a (t
kt -> t -> t
forall a. Num a => a -> a -> a
-t
1) (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
a) | t
a<-[t
1..(t -> t -> t
forall a. Ord a => a -> a -> a
min t
h (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
kt -> t -> t
forall a. Num a => a -> a -> a
+t
1))] ]

--------------------------------------------------------------------------------
-- Partitions with only odd\/distinct parts

{-
-- | Partitions of @n@ with only odd parts
partitionsWithOddParts :: Int -> [Partition]
partitionsWithOddParts d = map Partition (go d d) where
  go _  0  = [[]]
  go !h !n = [ a:as | a<-[1,3..min n h], as <- go a (n-a) ]
-}

{-
-- > length (partitionsWithDistinctParts d) == length (partitionsWithOddParts d)
--
partitionsWithDistinctParts :: Int -> [Partition]
partitionsWithDistinctParts d = map Partition (go d d) where
  go _  0  = [[]]
  go !h !n = [ a:as | a<-[1..min n h], as <- go (a-1) (n-a) ]
-}

--------------------------------------------------------------------------------