-- |
-- Module      : Data.Rhythm.Compositions
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
--
-- [Combinatorial compositions](https://mathworld.wolfram.com/Composition.html),
-- i.e., [partitions]("Data.Rhythm.Partitions") in which order is significant.
module Data.Rhythm.Compositions where

import Data.Bool (bool)
import Math.Combinat.Compositions (Composition, compositions1)
import System.Random (randomIO)

-- | All positive compositions of a given number.
--
-- >>> compositions 4
-- [[4],[1,3],[2,2],[3,1],[1,1,2],[1,2,1],[2,1,1],[1,1,1,1]]
compositions :: (Integral a) => a -> [Composition]
compositions :: forall a. Integral a => a -> [Composition]
compositions a
n = (a -> [Composition]) -> [a] -> [Composition]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (a -> a -> [Composition]
forall a. Integral a => a -> a -> [Composition]
`compositions1` a
n) [a
1 .. a
n]

-- | All positive compositions with allowed parts.
--
-- >>> compositionsAllowed [1,2] 4
-- [[2,2],[1,1,2],[1,2,1],[2,1,1],[1,1,1,1]]
compositionsAllowed :: (Integral a) => [Int] -> a -> [Composition]
compositionsAllowed :: forall a. Integral a => Composition -> a -> [Composition]
compositionsAllowed Composition
allowed = (Composition -> Bool) -> [Composition] -> [Composition]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Bool) -> Composition -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Composition -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Composition
allowed)) ([Composition] -> [Composition])
-> (a -> [Composition]) -> a -> [Composition]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [Composition]
forall a. Integral a => a -> [Composition]
compositions

-- | Positive compositions of a given length.
--
-- >>> compositionsLength 2 5
-- [[1,4],[2,3],[3,2],[4,1]]
--
-- The number of positive compositions of \(n\) into \(k\) parts is given by the
-- following formula.
--
-- \[
--   \begin{align*}
--     C_k(n) &= \binom{n - 1}{k - 1} \\
--     &= \frac{(n-1)!}{(k-1)!(n-k)!}
--   \end{align*}
-- \]
--
-- >>> let _C k n = toInteger (length (compositionsLength k n))
-- >>> let fact n = product [1 .. n]
-- >>> _C 2 5 == fact (5 - 1) `div` (fact (2 - 1) * fact (5 - 2))
-- True
compositionsLength :: Int -> Int -> [[Int]]
compositionsLength :: Int -> Int -> [Composition]
compositionsLength = Int -> Int -> [Composition]
forall a. Integral a => a -> a -> [Composition]
compositions1

-- | Positive compositions of a given length with allowed parts.
--
-- >>> compositionsLengthAllowed 2 [2,3] 5
-- [[2,3],[3,2]]
--
-- >>> filter (all (`elem` [2,3])) (compositionsLength 2 5)
-- [[2,3],[3,2]]
compositionsLengthAllowed :: Int -> [Int] -> Int -> [Composition]
compositionsLengthAllowed :: Int -> Composition -> Int -> [Composition]
compositionsLengthAllowed Int
len Composition
allowed =
  (Composition -> Bool) -> [Composition] -> [Composition]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Bool) -> Composition -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Composition -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Composition
allowed)) ([Composition] -> [Composition])
-> (Int -> [Composition]) -> Int -> [Composition]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> [Composition]
compositionsLength Int
len

-- | Generate a random positive composition of a given number.
--
-- >>> sum <$> randomComposition 13
-- 13
randomComposition :: Int -> IO Composition
randomComposition :: Int -> IO Composition
randomComposition Int
n = Int -> (Int, Composition) -> IO Composition
forall {f :: * -> *} {a}.
(Num a, MonadIO f) =>
Int -> (a, [a]) -> f [a]
go Int
1 (Int
1, [])
  where
    go :: Int -> (a, [a]) -> f [a]
go Int
i (a
p, [a]
acc)
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n = [a] -> f [a]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
p a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)
      | Bool
otherwise = Int -> (a, [a]) -> f [a]
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((a, [a]) -> f [a]) -> (Bool -> (a, [a])) -> Bool -> f [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, [a]) -> (a, [a]) -> Bool -> (a, [a])
forall a. a -> a -> Bool -> a
bool (a
1, a
p a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) (a
p a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, [a]
acc) (Bool -> f [a]) -> f Bool -> f [a]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< f Bool
forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO

-- | Generate a random positive composition of a given length.
--
-- >>> sum <$> randomCompositionLength 3 33
-- 33
randomCompositionLength :: Int -> Int -> IO Composition
randomCompositionLength :: Int -> Int -> IO Composition
randomCompositionLength Int
len Int
n = Int -> (Int, Int, Composition) -> IO Composition
forall {a} {m :: * -> *} {a}.
(Ord a, Num a, Num a, MonadIO m) =>
a -> (a, a, [a]) -> m [a]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
1, [])
  where
    go :: a -> (a, a, [a]) -> m [a]
go a
np (a
mp, a
j, [a]
acc)
      | a
mp a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = a -> (a, a, [a]) -> m [a]
go (a
np a -> a -> a
forall a. Num a => a -> a -> a
- a
1) ((a, a, [a]) -> m [a]) -> (Bool -> (a, a, [a])) -> Bool -> m [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a, [a]) -> (a, a, [a]) -> Bool -> (a, a, [a])
forall a. a -> a -> Bool -> a
bool (a
mp a -> a -> a
forall a. Num a => a -> a -> a
- a
1, a
1, a
j a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) (a
mp, a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
1, [a]
acc) (Bool -> m [a]) -> m Bool -> m [a]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Bool
forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO
      | Bool
otherwise = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
np a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)