-- | Compositions. 
--
-- See eg. <http://en.wikipedia.org/wiki/Composition_%28combinatorics%29>
--

module Math.Combinat.Compositions where

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

import System.Random

import Math.Combinat.Sets    ( randomChoice )
import Math.Combinat.Numbers ( factorial , binomial )
import Math.Combinat.Helper

--------------------------------------------------------------------------------
-- * generating all compositions

-- | A /composition/ of an integer @n@ into @k@ parts is an ordered @k@-tuple of nonnegative (sometimes positive) integers
-- whose sum is @n@.
type Composition = [Int]

-- | Compositions fitting into a given shape and having a given degree.
--   The order is lexicographic, that is, 
--
-- > sort cs == cs where cs = compositions' shape k
--
compositions'  
  :: [Int]         -- ^ shape
  -> Int           -- ^ sum
  -> [[Int]]
compositions' :: [Int] -> Int -> [[Int]]
compositions' [] Int
0 = [[]]
compositions' [] Int
_ = []
compositions' shape :: [Int]
shape@(Int
s:[Int]
ss) Int
n = 
  [ Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
xs | Int
x <- [Int
0..Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
s Int
n] , [Int]
xs <- [Int] -> Int -> [[Int]]
compositions' [Int]
ss (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
x) ] 

countCompositions' :: [Int] -> Int -> Integer
countCompositions' :: [Int] -> Int -> Integer
countCompositions' [] Int
0 = Integer
1
countCompositions' [] Int
_ = Integer
0
countCompositions' shape :: [Int]
shape@(Int
s:[Int]
ss) Int
n = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum 
  [ [Int] -> Int -> Integer
countCompositions' [Int]
ss (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
x) | Int
x <- [Int
0..Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
s Int
n] ] 

-- | All positive compositions of a given number (filtrated by the length). 
-- Total number of these is @2^(n-1)@
allCompositions1 :: Int -> [[Composition]]
allCompositions1 :: Int -> [[[Int]]]
allCompositions1 Int
n = (Int -> [[Int]]) -> [Int] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
d -> Int -> Int -> [[Int]]
forall a. Integral a => a -> a -> [[Int]]
compositions1 Int
d Int
n) [Int
1..Int
n] 

-- | All compositions fitting into a given shape.
allCompositions' :: [Int] -> [[Composition]]
allCompositions' :: [Int] -> [[[Int]]]
allCompositions' [Int]
shape = (Int -> [[Int]]) -> [Int] -> [[[Int]]]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> Int -> [[Int]]
compositions' [Int]
shape) [Int
0..Int
d] where d :: Int
d = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int]
shape

-- | Nonnegative compositions of a given length.
compositions 
  :: Integral a 
  => a       -- ^ length
  -> a       -- ^ sum
  -> [[Int]]
compositions :: a -> a -> [[Int]]
compositions a
len' a
d' = [Int] -> Int -> [[Int]]
compositions' (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
len Int
d) Int
d where
  len :: Int
len = a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
len'
  d :: Int
d   = a -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d'

-- | # = \\binom { len+d-1 } { len-1 }
countCompositions :: Integral a => a -> a -> Integer
countCompositions :: a -> a -> Integer
countCompositions a
len a
d = a -> a -> Integer
forall a. Integral a => a -> a -> Integer
binomial (a
lena -> a -> a
forall a. Num a => a -> a -> a
+a
da -> a -> a
forall a. Num a => a -> a -> a
-a
1) (a
lena -> a -> a
forall a. Num a => a -> a -> a
-a
1)

-- | Positive compositions of a given length.
compositions1  
  :: Integral a 
  => a       -- ^ length
  -> a       -- ^ sum
  -> [[Int]]
compositions1 :: a -> a -> [[Int]]
compositions1 a
len a
d 
  | a
len a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
d   = []
  | Bool
otherwise = ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> [Int]
plus1 ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ a -> a -> [[Int]]
forall a. Integral a => a -> a -> [[Int]]
compositions a
len (a
da -> a -> a
forall a. Num a => a -> a -> a
-a
len)
  where
    plus1 :: [Int] -> [Int]
plus1 = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    -- len = fromIntegral len'
    -- d   = fromIntegral d'

countCompositions1 :: Integral a => a -> a -> Integer
countCompositions1 :: a -> a -> Integer
countCompositions1 a
len a
d = a -> a -> Integer
forall a. Integral a => a -> a -> Integer
countCompositions a
len (a
da -> a -> a
forall a. Num a => a -> a -> a
-a
len)

--------------------------------------------------------------------------------
-- * random compositions

-- | @randomComposition k n@ returns a uniformly random composition 
-- of the number @n@ as an (ordered) sum of @k@ /nonnegative/ numbers
randomComposition :: RandomGen g => Int -> Int -> g -> ([Int],g)
randomComposition :: Int -> Int -> g -> ([Int], g)
randomComposition Int
k Int
n g
g0 = 
  if Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
1 Bool -> Bool -> Bool
|| Int
nInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0 
    then [Char] -> ([Int], g)
forall a. HasCallStack => [Char] -> a
error [Char]
"randomComposition: k should be positive, and n should be nonnegative" 
    else ([Int]
comp, g
g1) 
  where
    ([Int]
cs,g
g1) = Int -> Int -> g -> ([Int], g)
forall g. RandomGen g => Int -> Int -> g -> ([Int], g)
randomChoice (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) g
g0
    comp :: [Int]
comp = (Int -> Int -> Int) -> [Int] -> [Int]
forall a b. (a -> a -> b) -> [a] -> [b]
pairsWith (\Int
x Int
y -> Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
0 Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
cs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
k])
  
-- | @randomComposition1 k n@ returns a uniformly random composition 
-- of the number @n@ as an (ordered) sum of @k@ /positive/ numbers
randomComposition1 :: RandomGen g => Int -> Int -> g -> ([Int],g)
randomComposition1 :: Int -> Int -> g -> ([Int], g)
randomComposition1 Int
k Int
n g
g0 = 
  if Int
kInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
1 Bool -> Bool -> Bool
|| Int
nInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
k 
    then [Char] -> ([Int], g)
forall a. HasCallStack => [Char] -> a
error [Char]
"randomComposition1: we require 0 < k <= n" 
    else ([Int]
comp, g
g1) 
  where
    ([Int]
cs,g
g1) = Int -> Int -> g -> ([Int], g)
forall g. RandomGen g => Int -> Int -> g -> ([Int], g)
randomComposition Int
k (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k) g
g0 
    comp :: [Int]
comp = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Int]
cs

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