-- | Calculs de combinaisons.
module Htirage.Combin where

import Data.Bool
import Data.Eq (Eq(..))
import Data.Foldable (any, sum)
import Data.Int (Int)
import Data.List ((!!), length)
import Data.Ord (Ord(..))
import Prelude (Integral(..), Num(..), pred, undefined)

-- | @'nCk' n k@ retourne le nombre de combinaisons
-- de longueur 'k' d’un ensemble de longueur 'n'.
--
-- Computed using the formula:
-- @'nCk' n (k+1) == 'nCk' n (k-1) * (n-k+1) / k@
nCk :: Integral i => i -> i -> i
n`nCk`k | n<0||k<0||n<k = undefined
        | otherwise     = go 1 1
        where
        go i acc = if k' < i then acc else go (i+1) (acc * (n-i+1) `div` i)
        -- Use a symmetry to compute over smaller numbers,
        -- which is more efficient and safer
        k' = if n`div`2 < k then n-k else k

-- | @combinOfRank n k r@ retourne les indices de permutation
-- de la combinaison de 'k' entiers parmi @[1..n]@
-- au rang lexicographique 'r' dans @[0..'nCk' n k - 1]@.
-- 
-- Construit chaque choix de la combinaison en prenant le prochain plus grand
-- dont le successeur engendre un nombre de combinaisons
-- qui dépasse le rang restant à atteindre.
--
-- DOC: <http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf>, p.26
combinOfRank :: Integral i => i -> i -> i -> [i]
combinOfRank n k rk | rk<0||n`nCk`k<rk = undefined
                    | otherwise = for1K 1 1 rk
        where
        for1K i j r | i <  k    = uptoRank i j r
                    | i == k    = [j+r] -- because when i == k, nbCombs is always 1
                    | otherwise = []
        uptoRank i j r | nbCombs <- (n-j)`nCk`(k-i)
                       , nbCombs <= r = uptoRank i (j+1) (r-nbCombs)
                       | otherwise    = j : for1K (i+1) (j+1) r

-- | @rankOfCombin n ns@ retourne le rang lexicographique dans @[0..'nCk' n (length ns) - 1]@
-- de la combinaison 'ns' d’entiers parmi @[1..n]@.
--
-- WARNING: 'ns' doit être triée de manière ascendante.
--
-- Compte le nombre de combinaisons précédant celle de rang 'r'.
--
-- DOC: <http://www.site.uottawa.ca/~lucia/courses/5165-09/GenCombObj.pdf>, pp.24-25
--
-- @
-- 'rankOfCombin' n ('combinOfRank' n k r) == r
-- 'combinOfRank' n ('length' ns) ('rankOfCombin' n ns) == ns
-- @
rankOfCombin :: Integral i => i -> [i] -> i
rankOfCombin n ns | any (\x -> x<1||n<x) ns || n<k = undefined
                  | otherwise = for1K 1 0 0 ns
        where
        k = fromInteger (toInteger (length ns))
        for1K _ r _ []      = r
        for1K i r x1 (x:xs) = for1K (i+1) r' x xs
                where r' = r + sum [ (n-j)`nCk`(k-i)
                                   | j <- [x1+1..x-1]
                                   ]

-- | @permute ps xs@ remplace chaque élément de 'ps'
-- par l’élement qu’il indexe dans 'xs' entre @[1..'length' xs]@.
permute :: [Int] -> [a] -> [a]
permute ps xs = [xs !! pred p | p <- ps]