-- SPDX-License-Identifier: MPL-2.0

{- | Exact combinatorics. Indices and ranks are zero based; enumerations use
lexicographic order. Counts use unbounded 'Integer'. Enumeration APIs take a
maximum output count and reject the request before constructing that output.
Scalar recurrences accept orders through 4096 (quadratic tables through 512).
-}
module HMath.Combinatorics (
    permutationsCount,
    combinationsCount,
    multinomial,
    permutationsOf,
    combinationsOf,
    rankPermutation,
    unrankPermutation,
    rankCombination,
    unrankCombination,
    derangementCount,
    derangements,
    integerPartitions,
    setPartitions,
    stirlingFirst,
    stirlingSecond,
    bellNumber,
    eulerianNumber,
    lahNumber,
    catalanNumber,
    fibonacci,
    lucas,
    pascalTriangle,
    inclusionExclusion,
    subsets,
    cartesianProduct,
    multisetPermutations,
    multisetCount,
    linearRecurrence,
    generatingCoefficients,
    convolveSeries,
    finiteDifferences,
    finiteSum,
) where

import Data.List (delete, group, sort)
import HMath.Error (MathError (..))

order :: Int -> Either MathError ()
order :: Int -> Either MathError ()
order Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
4096 = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"combinatorial order must be in [0,4096]")
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

budget :: Int -> Integer -> Either MathError ()
budget :: Int -> Integer -> Either MathError ()
budget Int
b Integer
n
    | Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative output budget")
    | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
b = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"combinatorial output budget" Int
b)
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

fact :: Int -> Integer
fact :: Int -> Integer
fact Int
n = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Integer
1 .. Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n]

choose :: Int -> Int -> Integer
choose :: Int -> Int -> Integer
choose Int
n Int
k
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n = Integer
0
    | Bool
otherwise = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) .. Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n] Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
fact Int
r
  where
    r :: Int
r = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k)

{- | Given n,k in [0,4096], return n!/(n-k)! exactly, or zero when k>n. Invalid orders fail.
Example: permutationsCount 4 2 == Right 12.
-}
permutationsCount :: Int -> Int -> Either MathError Integer
permutationsCount :: Int -> Int -> Either MathError Integer
permutationsCount Int
n Int
k = do
    Int -> Either MathError ()
order Int
n
    Int -> Either MathError ()
order Int
k
    Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n then Integer
0 else Int -> Integer
fact Int
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
fact (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k))

{- | Given n,k in [0,4096], return the exact binomial coefficient, or zero when k>n. Invalid
orders fail. Example: combinationsCount 4 2 == Right 6.
-}
combinationsCount :: Int -> Int -> Either MathError Integer
combinationsCount :: Int -> Int -> Either MathError Integer
combinationsCount Int
n Int
k = do Int -> Either MathError ()
order Int
n; Int -> Either MathError ()
order Int
k; Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Integer
choose Int
n Int
k)

{- | Given nonnegative group sizes whose sum is \<=4096, return
factorial(total)/product(factorial sizes). Invalid sizes/total fail. Example:
multinomial [2,1] == Right 3.
-}
multinomial :: [Int] -> Either MathError Integer
multinomial :: [Int] -> Either MathError Integer
multinomial [Int]
ks = do
    (Int -> Either MathError ()) -> [Int] -> Either MathError ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Int -> Either MathError ()
order [Int]
ks
    let n :: Integer
n = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
ks)
    if Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
4096
        then MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"multinomial total exceeds 4096")
        else Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Integer
fact (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
n) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
fact [Int]
ks))

-- | Distinct positions, including when values compare equal.
permutationsOf :: Int -> Int -> [a] -> Either MathError [[a]]
permutationsOf :: forall a. Int -> Int -> [a] -> Either MathError [[a]]
permutationsOf Int
b Int
k [a]
xs = do
    n <- Int -> Int -> Either MathError Integer
permutationsCount ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) Int
k
    budget b n
    pure (perm k xs)
  where
    perm :: t -> [a] -> [[a]]
perm t
0 [a]
_ = [[]]
    perm t
j [a]
ys = [a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
zs | (Int
i, a
v) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [a]
ys, [a]
zs <- t -> [a] -> [[a]]
perm (t
j t -> t -> t
forall a. Num a => a -> a -> a
- t
1) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
ys [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
ys)]

{- | Given output budget, selection length k, then list, enumerate position-distinct
k-element combinations in source order. Invalid order or insufficient budget fails
before output construction.
-}
combinationsOf :: Int -> Int -> [a] -> Either MathError [[a]]
combinationsOf :: forall a. Int -> Int -> [a] -> Either MathError [[a]]
combinationsOf Int
b Int
k [a]
xs = do
    n <- Int -> Int -> Either MathError Integer
combinationsCount ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) Int
k; budget b n; pure (comb k xs)
  where
    comb :: t -> [a] -> [[a]]
comb t
0 [a]
_ = [[]]
    comb t
_ [] = []
    comb t
j (a
v : [a]
vs) = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) (t -> [a] -> [[a]]
comb (t
j t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [a]
vs) [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ t -> [a] -> [[a]]
comb t
j [a]
vs

-- | A permutation of @[0,n)@, in lexicographic Lehmer order.
rankPermutation :: [Int] -> Either MathError Integer
rankPermutation :: [Int] -> Either MathError Integer
rankPermutation [Int]
p = do
    Int -> Either MathError ()
order ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
p)
    if [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort [Int]
p [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        then MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"not a permutation")
        else Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int] -> [Int] -> Integer
forall {a}. Ord a => [a] -> [a] -> Integer
go [Int]
p [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
  where
    go :: [a] -> [a] -> Integer
go [] [a]
_ = Integer
0
    go (a
x : [a]
xs) [a]
ys = Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
x) [a]
ys)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
fact ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ [a] -> [a] -> Integer
go [a]
xs (a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
delete a
x [a]
ys)

{- | Given size n in [0,4096] then zero-based rank in [0,n!), return the lexicographic
permutation of [0..n-1]. Invalid size/rank fails.
-}
unrankPermutation :: Int -> Integer -> Either MathError [Int]
unrankPermutation :: Int -> Integer -> Either MathError [Int]
unrankPermutation Int
n Integer
r = do
    Int -> Either MathError ()
order Int
n
    if Integer
r Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
r Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Integer
fact Int
n
        then MathError -> Either MathError [Int]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"permutation rank out of bounds")
        else [Int] -> Either MathError [Int]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Integer -> [Int] -> [Int]
forall {a}. Eq a => Integer -> [a] -> [a]
go Integer
r [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
  where
    go :: Integer -> [a] -> [a]
go Integer
_ [] = []
    go Integer
q [a]
xs =
        let (Integer
i, Integer
t) = Integer
q Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int -> Integer
fact ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1); v :: a
v = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
i
         in a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Integer -> [a] -> [a]
go Integer
t (a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
delete a
v [a]
xs)

{- | Given universe size n then strictly increasing distinct indices in [0,n), return the
zero-based lexicographic combination rank. Invalid order/indices fail.
-}
rankCombination :: Int -> [Int] -> Either MathError Integer
rankCombination :: Int -> [Int] -> Either MathError Integer
rankCombination Int
n [Int]
xs = do
    Int -> Either MathError ()
order Int
n
    if [Int]
xs [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort [Int]
xs Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [[Int]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [[Int]]
forall a. Eq a => [a] -> [[a]]
group [Int]
xs) Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
x -> Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n) [Int]
xs
        then MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"combination must contain increasing indices")
        else Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> [Int] -> Integer
go Int
0 [Int]
xs)
  where
    go :: Int -> [Int] -> Integer
go Int
_ [] = Integer
0
    go Int
lo (Int
x : [Int]
rest) = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Int -> Integer
choose (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
rest) | Int
v <- [Int
lo .. Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> [Int] -> Integer
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
rest

{- | Given universe size n, selection size k, then zero-based rank, return the increasing
index list. Orders must be in [0,4096] and rank must be below choose(n,k); otherwise
fail.
-}
unrankCombination :: Int -> Int -> Integer -> Either MathError [Int]
unrankCombination :: Int -> Int -> Integer -> Either MathError [Int]
unrankCombination Int
n Int
k Integer
r = do
    Int -> Either MathError ()
order Int
n
    Int -> Either MathError ()
order Int
k
    if Integer
r Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
r Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Int -> Integer
choose Int
n Int
k
        then MathError -> Either MathError [Int]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"combination rank out of bounds")
        else [Int] -> Either MathError [Int]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Integer -> [Int]
go Int
0 Int
k Integer
r)
  where
    go :: Int -> Int -> Integer -> [Int]
go Int
_ Int
0 Integer
_ = []
    go Int
lo Int
j Integer
q = Int -> Integer -> [Int]
find Int
lo Integer
q
      where
        find :: Int -> Integer -> [Int]
find Int
x Integer
t
            | Integer
t Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
count = Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> Int -> Integer -> [Int]
go (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer
t
            | Bool
otherwise = Int -> Integer -> [Int]
find (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Integer
t Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
count)
          where
            count :: Integer
count = Int -> Int -> Integer
choose (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

{- | Return exact number of permutations of n elements having no fixed points. Require n in
[0,4096]; invalid orders fail. D0=1, D1=0. Example: derangementCount 3 == Right 2.
-}
derangementCount :: Int -> Either MathError Integer
derangementCount :: Int -> Either MathError Integer
derangementCount Int
n = do
    Int -> Either MathError ()
order Int
n
    Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst (((Integer, Integer) -> Int -> (Integer, Integer))
-> (Integer, Integer) -> [Int] -> (Integer, Integer)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(Integer
a, Integer
b) Int
k -> (Integer
b, Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b))) (Integer
1, Integer
0) [Int
2 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1]))

{- | Enumeration uses the factorial upper bound, so a request may be rejected
even if its filtered derangement output would fit.
-}
derangements :: Int -> Int -> Either MathError [[Int]]
derangements :: Int -> Int -> Either MathError [[Int]]
derangements Int
b Int
n = do
    ps <- Int -> Int -> [Int] -> Either MathError [[Int]]
forall a. Int -> Int -> [a] -> Either MathError [[a]]
permutationsOf Int
b Int
n [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    pure (filter (and . zipWith (/=) [0 ..]) ps)

{- | Nonincreasing positive parts. The budget counts search nodes as well as
outputs, preventing long searches with a tiny output limit.
-}
integerPartitions :: Int -> Int -> Either MathError [[Int]]
integerPartitions :: Int -> Int -> Either MathError [[Int]]
integerPartitions Int
b Int
n = do
    Int -> Either MathError ()
order Int
n
    if Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        then MathError -> Either MathError [[Int]]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"negative partition budget")
        else ([[Int]], Int) -> [[Int]]
forall a b. (a, b) -> a
fst (([[Int]], Int) -> [[Int]])
-> Either MathError ([[Int]], Int) -> Either MathError [[Int]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Int -> Either MathError ([[Int]], Int)
forall {b} {a}.
(Ord b, Ord a, Enum a, Num b, Num a) =>
b -> a -> a -> Either MathError ([[a]], b)
go Int
b Int
n Int
n
  where
    go :: b -> a -> a -> Either MathError ([[a]], b)
go b
fuel a
k a
hi
        | b
fuel b -> b -> Bool
forall a. Ord a => a -> a -> Bool
<= b
0 = MathError -> Either MathError ([[a]], b)
forall a b. a -> Either a b
Left (String -> Int -> MathError
NonConvergence String
"integer partitions" Int
b)
        | a
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = ([[a]], b) -> Either MathError ([[a]], b)
forall a b. b -> Either a b
Right ([[]], b
fuel b -> b -> b
forall a. Num a => a -> a -> a
- b
1)
        | Bool
otherwise = (Either MathError ([[a]], b) -> a -> Either MathError ([[a]], b))
-> Either MathError ([[a]], b)
-> [a]
-> Either MathError ([[a]], b)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Either MathError ([[a]], b) -> a -> Either MathError ([[a]], b)
step (([[a]], b) -> Either MathError ([[a]], b)
forall a b. b -> Either a b
Right ([], b
fuel b -> b -> b
forall a. Num a => a -> a -> a
- b
1)) [a -> a -> a
forall a. Ord a => a -> a -> a
min a
k a
hi, a -> a -> a
forall a. Ord a => a -> a -> a
min a
k a
hi a -> a -> a
forall a. Num a => a -> a -> a
- a
1 .. a
1]
      where
        step :: Either MathError ([[a]], b) -> a -> Either MathError ([[a]], b)
step Either MathError ([[a]], b)
prev a
x = do
            (out, f) <- Either MathError ([[a]], b)
prev
            (ys, f') <- go f (k - x) x
            pure (out ++ map (x :) ys, f')

-- | Set partitions of @[0,n)@, using restricted-growth order of block labels.
setPartitions :: Int -> Int -> Either MathError [[[Int]]]
setPartitions :: Int -> Int -> Either MathError [[[Int]]]
setPartitions Int
b Int
n = do
    Int -> Either MathError ()
order Int
n
    count <- Int -> Either MathError Integer
bellNumber Int
n
    budget b count
    pure
        [ [[i | (i, v) <- zip [0 ..] labels, v == j] | j <- [0 .. maximum (-1 : labels)]]
        | labels <- grow n [] (-1 :: Int)
        ]
  where
    grow :: t -> [t] -> t -> [[t]]
grow t
0 [t]
acc t
_ = [[t]
acc]
    grow t
k [t]
acc t
hi = [[[t]]] -> [[t]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [t -> [t] -> t -> [[t]]
grow (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1) ([t]
acc [t] -> [t] -> [t]
forall a. [a] -> [a] -> [a]
++ [t
j]) (t -> t -> t
forall a. Ord a => a -> a -> a
max t
hi t
j) | t
j <- [t
0 .. t
hi t -> t -> t
forall a. Num a => a -> a -> a
+ t
1]]

triangle :: (Int -> Int -> Integer -> Integer -> Integer) -> Int -> Int -> Either MathError Integer
triangle :: (Int -> Int -> Integer -> Integer -> Integer)
-> Int -> Int -> Either MathError Integer
triangle Int -> Int -> Integer -> Integer -> Integer
f Int
n Int
k = do
    Int -> Either MathError ()
order Int
n
    Int -> Either MathError ()
order Int
k
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
512
        then MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"triangular recurrence order exceeds 512")
        else Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((([Integer] -> [Integer]) -> [Integer] -> [[Integer]]
forall a. (a -> a) -> a -> [a]
iterate [Integer] -> [Integer]
next [Integer
1] [[Integer]] -> Int -> [Integer]
forall a. HasCallStack => [a] -> Int -> a
!! Int
n) [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n then Integer
0 else Integer
1)
  where
    next :: [Integer] -> [Integer]
next [Integer]
row = let i :: Int
i = [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
row in [Int -> Int -> Integer -> Integer -> Integer
f Int
i Int
j ([Integer] -> Int -> Integer
forall {a}. Num a => [a] -> Int -> a
at [Integer]
row (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) ([Integer] -> Int -> Integer
forall {a}. Num a => [a] -> Int -> a
at [Integer]
row Int
j) | Int
j <- [Int
0 .. Int
i]]
    at :: [a] -> Int -> a
at [a]
xs Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs = a
0 | Bool
otherwise = [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i

-- | Signed first-kind Stirling numbers, coefficients of the falling factorial.
stirlingFirst :: Int -> Int -> Either MathError Integer
stirlingFirst :: Int -> Int -> Either MathError Integer
stirlingFirst = (Int -> Int -> Integer -> Integer -> Integer)
-> Int -> Int -> Either MathError Integer
triangle (\Int
n Int
_ Integer
a Integer
b -> Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b)

{- | Given n,k, return the exact number of partitions into k nonempty unlabeled blocks.
Require n\<=512 and nonnegative bounded orders; k>n gives zero.
-}
stirlingSecond :: Int -> Int -> Either MathError Integer
stirlingSecond :: Int -> Int -> Either MathError Integer
stirlingSecond = (Int -> Int -> Integer -> Integer -> Integer)
-> Int -> Int -> Either MathError Integer
triangle (\Int
_ Int
k Integer
a Integer
b -> Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b)

{- | Return the exact number of set partitions of n labeled objects, summing second-kind
Stirling numbers. The triangular recurrence requires n in [0,512].
-}
bellNumber :: Int -> Either MathError Integer
bellNumber :: Int -> Either MathError Integer
bellNumber Int
n = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Integer] -> Integer)
-> Either MathError [Integer] -> Either MathError Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> Either MathError Integer)
-> [Int] -> Either MathError [Integer]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Int -> Either MathError Integer
stirlingSecond Int
n) [Int
0 .. Int
n]

{- | Given n,k, return the Eulerian recurrence counting permutations by ascents. Require
n\<=512 and bounded nonnegative orders; k>n gives zero.
-}
eulerianNumber :: Int -> Int -> Either MathError Integer
eulerianNumber :: Int -> Int -> Either MathError Integer
eulerianNumber = (Int -> Int -> Integer -> Integer -> Integer)
-> Int -> Int -> Either MathError Integer
triangle (\Int
n Int
k Integer
a Integer
b -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
b)

{- | Given n,k in [0,4096], return unsigned Lah number n!/k!*choose(n-1,k-1), with (0,0)=1
and impossible positive-size partitions giving zero. Invalid orders fail.
-}
lahNumber :: Int -> Int -> Either MathError Integer
lahNumber :: Int -> Int -> Either MathError Integer
lahNumber Int
n Int
k = do
    Int -> Either MathError ()
order Int
n
    Int -> Either MathError ()
order Int
k
    Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        ( if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then Integer
1
            else if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n then Integer
0 else Int -> Int -> Integer
choose (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
fact Int
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
fact Int
k
        )

{- | Return exact Catalan number choose(2*n,n)/(n+1), requiring n in [0,4096]. Example:
catalanNumber 3 == Right 5.
-}
catalanNumber :: Int -> Either MathError Integer
catalanNumber :: Int -> Either MathError Integer
catalanNumber Int
n = do Int -> Either MathError ()
order Int
n; Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Integer
choose (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Int
n Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))

{- | Return exact zero-based Fibonacci number by fast doubling for n in [0,4096]. Invalid
index fails. Example: fibonacci 10 == Right 55.
-}
fibonacci :: Int -> Either MathError Integer
fibonacci :: Int -> Either MathError Integer
fibonacci Int
n = do Int -> Either MathError ()
order Int
n; Integer -> Either MathError Integer
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst (Int -> (Integer, Integer)
forall {a} {b}. (Integral a, Num b) => a -> (b, b)
pair Int
n))
  where
    pair :: a -> (b, b)
pair a
0 = (b
0, b
1)
    pair a
k =
        let (b
a, b
b) = a -> (b, b)
pair (a
k a -> a -> a
forall a. Integral a => a -> a -> a
`div` a
2); c :: b
c = b
a b -> b -> b
forall a. Num a => a -> a -> a
* (b
2 b -> b -> b
forall a. Num a => a -> a -> a
* b
b b -> b -> b
forall a. Num a => a -> a -> a
- b
a); d :: b
d = b
a b -> b -> b
forall a. Num a => a -> a -> a
* b
a b -> b -> b
forall a. Num a => a -> a -> a
+ b
b b -> b -> b
forall a. Num a => a -> a -> a
* b
b
         in if a -> Bool
forall a. Integral a => a -> Bool
even a
k then (b
c, b
d) else (b
d, b
c b -> b -> b
forall a. Num a => a -> a -> a
+ b
d)

{- | Return zero-based Lucas number using neighboring Fibonacci numbers. n and n+1 must both
satisfy the [0,4096] bound, so n\<=4095. Invalid index fails.
-}
lucas :: Int -> Either MathError Integer
lucas :: Int -> Either MathError Integer
lucas Int
n = do a <- Int -> Either MathError Integer
fibonacci Int
n; b <- fibonacci (n + 1); pure (2 * b - a)

{- | Return Pascal rows zero through n inclusive, starting [[1]]. n must lie in [0,512];
invalid orders fail.
-}
pascalTriangle :: Int -> Either MathError [[Integer]]
pascalTriangle :: Int -> Either MathError [[Integer]]
pascalTriangle Int
n = do
    Int -> Either MathError ()
order Int
n
    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
512
        then MathError -> Either MathError [[Integer]]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"Pascal order exceeds 512")
        else [[Integer]] -> Either MathError [[Integer]]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> [[Integer]] -> [[Integer]]
forall a. Int -> [a] -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (([Integer] -> [Integer]) -> [Integer] -> [[Integer]]
forall a. (a -> a) -> a -> [a]
iterate (\[Integer]
r -> (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
0 Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: [Integer]
r) ([Integer]
r [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [Integer
0])) [Integer
1]))

{- | Union cardinality from all nonempty intersection sizes, indexed by subset
bit mask 1 through 2^n-1. The caller supplies consistent intersection counts.
-}
inclusionExclusion :: Int -> [Integer] -> Either MathError Integer
inclusionExclusion :: Int -> [Integer] -> Either MathError Integer
inclusionExclusion Int
n [Integer]
sizes
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
20 Bool -> Bool -> Bool
|| [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
sizes Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
2 Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Bool -> Bool -> Bool
|| (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0) [Integer]
sizes =
        MathError -> Either MathError Integer
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"invalid intersection table")
    | Bool
otherwise = Integer -> Either MathError Integer
forall a b. b -> Either a b
Right ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [if Int -> Bool
forall a. Integral a => a -> Bool
odd (Int -> Int
ones Int
i) then Integer
x else -Integer
x | (Int
i, Integer
x) <- [Int] -> [Integer] -> [(Int, Integer)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1 :: Int ..] [Integer]
sizes])
  where
    ones :: Int -> Int
ones Int
0 = Int
0 :: Int; ones Int
k = Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int
ones (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

{- | Given maximum output count then list, enumerate its position-based powerset. Reject
input order >4096 or a budget smaller than 2^length. Equal input values can yield
equal-valued subsets.
-}
subsets :: Int -> [a] -> Either MathError [[a]]
subsets :: forall a. Int -> [a] -> Either MathError [[a]]
subsets Int
b [a]
xs = do
    Int -> Either MathError ()
order ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
    Int -> Integer -> Either MathError ()
budget Int
b (Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
    [[a]] -> Either MathError [[a]]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> [[a]] -> [[a]]) -> [[a]] -> [a] -> [[a]]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x [[a]]
ps -> [[a]]
ps [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) [[a]]
ps) [[]] [a]
xs)

{- | Given output budget then a list of lists, return all ordered choices using sequence.
Reject negative/insufficient budget before constructing results. Empty outer list
yields one empty choice.
-}
cartesianProduct :: Int -> [[a]] -> Either MathError [[a]]
cartesianProduct :: forall a. Int -> [[a]] -> Either MathError [[a]]
cartesianProduct Int
b [[a]]
xs = do
    Int -> Integer -> Either MathError ()
budget Int
b ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product (([a] -> Integer) -> [[a]] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> ([a] -> Int) -> [a] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[a]]
xs)); [[a]] -> Either MathError [[a]]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[a]] -> [[a]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [[a]]
xs)

{- | Return the multinomial count of distinct permutations of a list, grouping equal values
in sorted order. Group sizes and total must satisfy the 4096-element combinatorial
bound.
-}
multisetCount :: (Ord a) => [a] -> Either MathError Integer
multisetCount :: forall a. Ord a => [a] -> Either MathError Integer
multisetCount [a]
xs = [Int] -> Either MathError Integer
multinomial (([a] -> Int) -> [[a]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group ([a] -> [a]
forall a. Ord a => [a] -> [a]
sort [a]
xs)))

{- | Given budget then ordered values, enumerate distinct multiset permutations in ascending
lexicographic order. Invalid size or insufficient output budget fails before
enumeration.
-}
multisetPermutations :: (Ord a) => Int -> [a] -> Either MathError [[a]]
multisetPermutations :: forall a. Ord a => Int -> [a] -> Either MathError [[a]]
multisetPermutations Int
b [a]
xs = do
    count <- [a] -> Either MathError Integer
forall a. Ord a => [a] -> Either MathError Integer
multisetCount [a]
xs; budget b count; pure (go (sort xs))
  where
    go :: [a] -> [[a]]
go [] = [[]]
    go [a]
ys = [a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
zs | a
v <- [a
v | a
v : [a]
_ <- [a] -> [[a]]
forall a. Eq a => [a] -> [[a]]
group [a]
ys], [a]
zs <- [a] -> [[a]]
go (a -> [a] -> [a]
forall a. Eq a => a -> [a] -> [a]
delete a
v [a]
ys)]

{- | @a[n] = sum c[i]*a[n-i-1]@; initial values have the same length as
coefficients. A finite exact prefix of at most 4096 coefficients is returned.
-}
linearRecurrence :: Int -> [Integer] -> [Integer] -> Either MathError [Integer]
linearRecurrence :: Int -> [Integer] -> [Integer] -> Either MathError [Integer]
linearRecurrence Int
n [Integer]
cs [Integer]
initial = do
    Int -> Either MathError ()
order Int
n
    if [Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Integer]
cs Bool -> Bool -> Bool
|| [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
cs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
initial
        then MathError -> Either MathError [Integer]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"recurrence initial state mismatch")
        else [Integer] -> Either MathError [Integer]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take Int
n ([Integer] -> [Integer]
go [Integer]
initial))
  where
    k :: Int
k = [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
cs
    go :: [Integer] -> [Integer]
go [Integer]
ys = [Integer]
ys [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [Integer] -> [Integer]
extend [Integer]
ys
    extend :: [Integer] -> [Integer]
extend [Integer]
ys = let z :: Integer
z = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((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]
cs ([Integer] -> [Integer]
forall a. [a] -> [a]
reverse (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
drop ([Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
ys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) [Integer]
ys))) in Integer
z Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: [Integer] -> [Integer]
extend (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
drop Int
1 [Integer]
ys [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [Integer
z])

{- | Formal power-series division @P(x)/Q(x)@ around zero, exact rationals.
The denominator constant must be nonzero; no analytic convergence is claimed.
-}
generatingCoefficients :: Int -> [Rational] -> [Rational] -> Either MathError [Rational]
generatingCoefficients :: Int -> [Rational] -> [Rational] -> Either MathError [Rational]
generatingCoefficients Int
n [Rational]
ps [Rational]
qs = do
    Int -> Either MathError ()
order Int
n
    if [Rational] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Rational]
qs Bool -> Bool -> Bool
|| [Rational] -> Int -> Rational
forall {a}. Num a => [a] -> Int -> a
get [Rational]
qs Int
0 Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0
        then MathError -> Either MathError [Rational]
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"series denominator has zero constant")
        else [Rational] -> Either MathError [Rational]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Rational] -> Int -> [Rational]
go [] Int
0)
  where
    get :: [a] -> Int -> a
get [a]
xs Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs then [a]
xs [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
i else a
0
    go :: [Rational] -> Int -> [Rational]
go [Rational]
acc Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = [Rational]
acc
        | Bool
otherwise =
            let v :: Rational
v = ([Rational] -> Int -> Rational
forall {a}. Num a => [a] -> Int -> a
get [Rational]
ps Int
i Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
- [Rational] -> Rational
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Rational] -> Int -> Rational
forall {a}. Num a => [a] -> Int -> a
get [Rational]
qs Int
j Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
* ([Rational]
acc [Rational] -> Int -> Rational
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j)) | Int
j <- [Int
1 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
i ([Rational] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Rational]
qs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]]) Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ [Rational] -> Int -> Rational
forall {a}. Num a => [a] -> Int -> a
get [Rational]
qs Int
0
             in [Rational] -> Int -> [Rational]
go ([Rational]
acc [Rational] -> [Rational] -> [Rational]
forall a. [a] -> [a] -> [a]
++ [Rational
v]) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{- | Given output length n then two coefficient lists, return the first n exact Integer
convolution coefficients, treating absent coefficients as zero. n must lie in [0,4096].
-}
convolveSeries :: Int -> [Integer] -> [Integer] -> Either MathError [Integer]
convolveSeries :: Int -> [Integer] -> [Integer] -> Either MathError [Integer]
convolveSeries Int
n [Integer]
a [Integer]
b = do
    Int -> Either MathError ()
order Int
n
    [Integer] -> Either MathError [Integer]
forall a. a -> Either MathError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        [[Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [[Integer]
a [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! Int
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* [Integer]
b [Integer] -> Int -> Integer
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j) | Int
j <- [Int
0 .. Int
i], Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
a, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Integer] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
b] | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

{- | Return successive forward-difference rows beginning with the original list and ending
at the last singleton. Empty input yields no rows. Example: finiteDifferences [1,4,9]
== [[1,4,9],[3,5],[2]].
-}
finiteDifferences :: (Num a) => [a] -> [[a]]
finiteDifferences :: forall a. Num a => [a] -> [[a]]
finiteDifferences [] = []
finiteDifferences [a]
xs = [a]
xs [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
forall a. Num a => [a] -> [[a]]
finiteDifferences ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1 [a]
xs) [a]
xs)

{- | Sum a finite list using its Num arithmetic. Empty input returns zero. Example:
finiteSum [1,2,3] == 6.
-}
finiteSum :: (Num a) => [a] -> a
finiteSum :: forall a. Num a => [a] -> a
finiteSum = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum