{-# LANGUAGE TupleSections #-}
module Multiset where
import qualified Data.Map.Strict as M
import qualified IntMapHelp as IM
import qualified Data.IntSet as IS
import Control.Monad

-- naive solution
powlistUpToN' :: Int -> [a] ->[[a]]
powlistUpToN' _ [] = [[]]
powlistUpToN' n (x:xs) 
    | n <=0     = [[]]
    | otherwise = ((x:) <$> powlistUpToN' (n-1) xs) ++ powlistUpToN' n xs

-- adapted from
-- https://stackoverflow.com/questions/21265454/subsequences-of-length-n-from-list-performance/59932616#59932616
-- uses dynamic programming: the important part is the use of "next" twice
-- there is a (probably) faster SO answer that produces power lists of size
-- exactly N, but that answer is harder to adapt
powlistUpToN :: Int -> [a] -> [[a]]
powlistUpToN n xs = concat $ drop (length xs-n) (subseqsBySize xs)
  where
    subseqsBySize [] = [[[]]]
    subseqsBySize (y:ys) =
        let next = subseqsBySize ys
        in zipWith (++) ([]:next) (map (map (y:)) next ++ [[]])
powlistUpToN'' :: Int -> [a] -> [[a]]
powlistUpToN'' n xs = 
    let l = length xs
--    in if n > l then [] else concat $ drop (l-n) (subseqsBySize xs)
    in if n > l 
        then concat $ subseqsBySize xs
        else concat $ drop (l-n) (subseqsBySize xs)
  where
    subseqsBySize [] = [[[]]]
    subseqsBySize (y:ys) =
        let next = subseqsBySize ys
        in zipWith (++) ([]:next) (map (map (y:)) next ++ [[]])
-- this is the code producing all exactly n length sublists
combinationsOf :: Int -> [a] -> [[a]]
combinationsOf 1 as        = map pure as
combinationsOf k' as@(_:xs) = run (l-1) (k'-1) as $ combinationsOf (k'-1) xs
  where
    l = length as
    run :: Int -> Int -> [a] -> [[a]] -> [[a]]
    run n k ys cs 
        | n == k    = map (ys ++) cs
        | otherwise = map (q:) cs ++ run (n-1) k qs (drop dc cs)
      where
        (q:qs) = take (n-k+1) ys
        dc     = product [(n-k+1)..(n-1)] `div` product [1..(k-1)]
combinationsOf _ [] = []

-- exponential, so don't use it on long lists
powlist :: [a] -> [[a]]
powlist = filterM (const [True,False])

toMultiset :: Ord a => [a] -> M.Map a Int
toMultiset = foldr (uncurry (M.insertWith (+)) . (, 1)) M.empty

invertIntMap :: Ord a => IM.IntMap a -> M.Map a IS.IntSet
invertIntMap = IM.foldrWithKey 
    (\k x -> M.insertWith IS.union x (IS.singleton k))
    M.empty