{-| Implementation of a heap which allows storing multiple values under one 
    priority (key), preserving the insertion order. It is based on Data.GHeap and Data.Sequence and virtually
    just interfacing the necessary functions.
-}
module Data.BagHeap (


    BagHeap(..),
    empty, isEmpty,
    
    singleton, insert,
    singletonL, insertL, 
    merge, 
    getMin, delMin, getDelMin,
    
    toList, prios, elems,
    fromList,

)where


import Data.GHeap (Heap)
import qualified Data.GHeap as H
import Data.Sequence (Seq, (><))
import qualified Data.Sequence as S
import qualified Data.Foldable as F
-- import qualified Data.List as L (null)
import Data.List (foldl', groupBy, partition,null)

import Data.Function (on)
import Control.Arrow

-- | A heap storing values with the same priority (key) 'p' in a sequence. 
data BagHeap p v = BH {unBH :: (Heap p (Seq v))} 


-- | Returns an empty heap
empty :: (Ord p) => BagHeap p v
empty = BH H.empty

-- | Returns 'True' if the heap is empty, 'alse' otherwise
isEmpty :: (Ord p) => BagHeap p v -> Bool
isEmpty = H.isEmpty . unBH

-- | creates a singleton heap from a priority\/value pair
singleton :: (Ord p) => p -> v -> BagHeap p v
singleton p v = BH $ H.singleton p (S.singleton v)  

-- | creates a singleton heap from a priority\/value-list pair. If 'lv' is 
--   empty it, the empty heap is returned
singletonL :: (Ord p) => p -> [v] -> BagHeap p v
singletonL p vl 
    | null vl   = empty 
    | otherwise = BH $ H.singleton p (S.fromList vl) 

-- | 'insert k v h' inserts the value 'k' with priority 'k' into heap 'h'
insert :: (Ord p) => p -> v -> BagHeap p v -> BagHeap p v
insert  = (flip merge . ) . singleton

-- | 'insertL k lv h' inserts the list of values 'lv' with priority 'k' into 
--   heap 'h'. If 'lv' is empty it, the empty heap is returned
insertL :: (Ord p) => p -> [v] -> BagHeap p v -> BagHeap p v
insertL  p vl 
    | null vl   = id
    | otherwise = flip merge (singletonL p vl)

-- | Merge two heaps into one
merge :: (Ord p) => BagHeap p v -> BagHeap p v -> BagHeap p v
merge = (BH .) . (on (H.mergeWith (><)) unBH)

--  | Returns the minimal priority\/value pair and the heap without it
getDelMin :: (Ord p) => BagHeap p v -> ((p,[v]),BagHeap p v)
getDelMin = ((id *** F.toList) *** BH) . H.getDelMin . unBH

----  | Returns the minimal priority\/value pair and the heap without it
--getDelMinWhich :: (Ord p) => (v -> Bool) ->  BagHeap p v -> ((p,[v]),BagHeap p v)
--getDelMinWhich f h = 
--    let ((p,(vt,vf)),h') = ((id *** ((partition f).F.toList)) *** BH) . H.getDelMin $ unBH  h
--    in ((p,vt), merge (singletonL p vf) h') 

-- | Returns the th minimal priority\/value pair
getMin :: (Ord p) => BagHeap p v -> (p,[v])
getMin = fst . getDelMin

-- | Returns the resulting heap after removing the minimal priority\/value pair
delMin :: (Ord p) => BagHeap p v -> BagHeap p v
delMin = snd . getDelMin

toList :: (Ord p) => BagHeap p v -> [(p,[v])]
toList = (map (id *** F.toList)) . H.toList . unBH

prios :: (Ord p) => BagHeap p v -> [p]
prios m = [p | (p,_) <- toList m]

elems :: (Ord p) => BagHeap p v -> [[v]]
elems m = [v | (_,v) <- toList m]

fromList :: (Ord p) => [(p,v)] -> BagHeap p v
fromList = (foldl' (flip $ uncurry insert) empty)

fromAscList :: (Ord p) => [(p,v)] -> BagHeap p v
fromAscList = (foldl' (flip $ uncurry insertL) empty) . groupAsc

groupAsc :: (Eq p) => [(p,v)] -> [(p,[v])]
groupAsc [] = []
groupAsc (x:xs) = cons x (groupAsc xs)
    where
    cons (p,v) [] = [(p,[v])]
    cons (p,v) xs@((p',vs):xss)
        | p == p'   = ((p',v:vs):xss)
        | otherwise = (p,[v]):xs 