{-# LANGUAGE TypeSynonymInstances #-}
module SmartGroup (Splittable, groupAll, groupNum, groupLog) where
import Data.Heap as Heap
import Data.Set as Set
import Data.Map as Map
import Data.Monoid
import Data.Char
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Ord

class Ord a => Splittable a where wordsOf :: a -> [a]
instance Splittable String where wordsOf = Prelude.filter ((>3) . length) . words
instance Splittable L.ByteString where
         wordsOf = Prelude.filter ((>3) . L.length) . L.splitWith isSpace
instance Splittable S.ByteString where
         wordsOf = Prelude.filter ((>3) . S.length) . S.splitWith isSpace

data StringL a = StringL {str :: a, count :: Int} deriving (Eq, Show)
instance Ord a => Ord (StringL a) where compare (StringL i a) (StringL x b) = compare a b `mappend` compare i x
data SizeMap s a = Unsplittable (Set a) | Splittable (Map (StringL s) (Set a)) deriving (Eq, Show)
instance (Ord s, Ord a) => Ord (SizeMap s a) where
         compare (Splittable a) (Splittable b) = compare (Map.size a) (Map.size b)
         compare (Unsplittable a) (Unsplittable b) = compare a b
         compare (Unsplittable _) (Splittable _) = LT
         compare (Splittable _) (Unsplittable _) = GT
type WordAssoc s a = MaxHeap (SizeMap s a)

toSet (Unsplittable x) = x
toSet (Splittable x) = Set.unions (Map.elems x)

intLog :: Int -> Int
intLog = truncate . logBase 2 . fromIntegral

groupWith :: (Ord a, Splittable s) => (Int -> WordAssoc s a -> WordAssoc s a) -> Int -> (a -> s) -> [a] -> [[a]]
groupWith f i c = mkList . f i . mkAssoc . mkMap c

-- | Divide list into as many groups as possible
groupAll :: (Ord a, Splittable s) => Int -> (a -> s) -> [a] -> [[a]]
groupAll = groupWith $ \i x->
         let cycleSplit n = case splitIt i n of
                  (Just a) -> cycleSplit a
                  Nothing -> n
         in cycleSplit x

-- | Divide list into about n different groups
groupNum :: (Ord a, Splittable s) => Int -> Int -> (a -> s) -> [a] -> [[a]]
groupNum i = groupWith $ \x->
         let splitTimes n m =
                if Heap.size m >= n then m else
                   case splitIt x m of
                     (Just a) -> splitTimes n a
                     Nothing -> m
         in splitTimes i

-- | Divide list into groups such that the amount of groups
--   equals the log of the number of elements
groupLog :: (Ord a, Splittable s) => Int -> (a -> s) -> [a] -> [[a]]
groupLog i f s = groupNum (intLog (length s)) i f s

mkMap :: (Ord a, Splittable s) => (a -> s) -> [a] -> Map s (Set a)
mkMap f = foldl (\m x->
        foldl (\m' i-> Map.alter (Just . maybe (Set.singleton x) (Set.insert x)) i m')
        m (wordsOf $ f x)) Map.empty

mkAssoc :: (Ord a, Splittable s) => Map s (Set a) -> MaxHeap (SizeMap s a)
mkAssoc m = Heap.singleton . Splittable . Map.mapKeys (\k-> StringL k (Set.size (m Map.! k))) $ m

splitIt :: (Ord a, Splittable s) => Int -> WordAssoc s a -> Maybe (WordAssoc s a)
splitIt i s = case Heap.view s of
         Nothing -> Nothing
         (Just ((Unsplittable _),_)) -> Nothing
         (Just ((Splittable x),xs)) -> do
               (as,b) <- Map.maxView x
               x1 <- sizeMap i $ flip Map.mapMaybe b $ \n->
                      let m = Set.difference n as
                      in if Set.null m then Nothing else Just m
               x2 <- sizeMap i $ flip Map.mapMaybe b $ \n->
                      let m = Set.intersection n as
                      in if Set.null m then Nothing else Just m
               let x3 = as Set.\\ (Set.unions (Map.elems b))
               return $ (if Set.null x3 then id else Heap.insert (Unsplittable x3)) $
                   Heap.insert x1 $ Heap.insert x2 xs

sizeMap :: (Ord a, Splittable s) => Int -> (Map (StringL s) (Set a)) -> Maybe (SizeMap s a)
sizeMap i m = if Map.null m then Nothing else Just $ case Map.findMax m of
        (StringL _ c,_) -> if c >= i then Splittable m
                 else Unsplittable (Set.unions (Map.elems m))

-- findMax fails on empty maps
-- maxView might not return a Just value
-- known offenders
   -- ["","","",""] (empty strings have no words to key off of)

mkList :: (Ord a, Splittable s) => WordAssoc s a -> [[a]]
mkList = Prelude.map (Set.toList . toSet) . Heap.toList