{-# LANGUAGE ScopedTypeVariables #-}
module Sums where
import Data.IntMap.Lazy (IntMap)
import Data.Maybe
import Data.Monoid ((<>))
import qualified Data.IntMap as IntMap

-- find all subsets which add upto the given size, given a function to
-- extract the size.  Only pick at most one element from each input subset.
sums :: forall a.(a -> Int) -> Int -> [[a]] -> [[a]]
sums size sum_ = fromMaybe [] . IntMap.lookup sum_ . sumsAux2
  where
    sumsAux2 :: [[a]] -> IntMap [[a]]
    sumsAux2 = foldr (\xs im -> foldr (sumsAux im) im xs) IntMap.empty

    sumsAux :: IntMap [[a]] -> a -> IntMap [[a]] -> IntMap [[a]]
    sumsAux origIm x newIm
      | size x > sum_ = newIm
      | otherwise =
          IntMap.insertWith (<>) (size x) [[x]] $
          IntMap.foldrWithKey (\key val im2 ->
                                 if key + size x > sum_ then im2
                                 else IntMap.insertWith (<>) (key + size x)
                                      (map (x:) val) im2)
          newIm origIm