{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Text.PageIO.LabelMap where

import StringTable.Atom
import qualified Data.IntMap as IM
import Data.Monoid

newtype LabelMap a = MkLabelMap { labelMap :: IM.IntMap a }
    deriving (Eq, Ord, Functor, Monoid)

newtype Label = MkLabel { labelAtom :: Atom }
    deriving (Eq, Ord)

instance Show a => Show (LabelMap a) where
    show = show . toList

instance Show Label where
    show = fromLabel

toLabel :: ToAtom a => a -> Label
toLabel = MkLabel . toAtom

fromLabel :: FromAtom a => Label -> a
fromLabel = fromAtom . labelAtom

fromList :: [(Label, a)] -> LabelMap a
fromList ps = MkLabelMap $ IM.fromList [ (fromLabel l, x) | (l, x) <- ps ]

fromListWith :: (a -> a -> a) -> [(Label, a)] -> LabelMap a
fromListWith f ps = MkLabelMap $ IM.fromListWith f [ (fromLabel l, x) | (l, x) <- ps ]

toList :: LabelMap a -> [(Label, a)]
toList lm = [ (keyToLabel l, x) | (l, x) <- IM.toList $ labelMap lm ]

{-# INLINE keyToLabel #-}
keyToLabel :: IM.Key -> Label
keyToLabel = MkLabel . unsafeIntToAtom

mapWithKey :: (Label -> a -> b) -> LabelMap a -> LabelMap b
mapWithKey f lm = MkLabelMap $ IM.mapWithKey (f . keyToLabel) (labelMap lm)

elems :: LabelMap a -> [a]
elems = IM.elems . labelMap

union :: LabelMap a -> LabelMap a -> LabelMap a
union (MkLabelMap x) (MkLabelMap y) = MkLabelMap (IM.union x y)

unionWith :: (a -> a -> a) -> LabelMap a -> LabelMap a -> LabelMap a
unionWith f (MkLabelMap x) (MkLabelMap y) = MkLabelMap (IM.unionWith f x y)

unions :: [LabelMap a] -> LabelMap a
unions = MkLabelMap . IM.unions . Prelude.map labelMap

unionsWith :: (a -> a -> a) -> [LabelMap a] -> LabelMap a
unionsWith f = MkLabelMap . IM.unionsWith f . Prelude.map labelMap

lookup :: Label -> LabelMap a -> Maybe a
lookup l = IM.lookup (fromLabel l) . labelMap

insert :: Label -> a -> LabelMap a -> LabelMap a
insert l v = MkLabelMap . IM.insert (fromLabel l) v . labelMap

insertWith :: (a -> a -> a) -> Label -> a -> LabelMap a -> LabelMap a
insertWith f l v = MkLabelMap . IM.insertWith f (fromLabel l) v . labelMap

member :: Label -> LabelMap a -> Bool
member l = IM.member (fromLabel l) . labelMap

keys :: LabelMap a -> [Label]
keys = Prelude.map keyToLabel . IM.keys . labelMap

filter :: (a -> Bool) -> LabelMap a -> LabelMap a
filter f = MkLabelMap . IM.filter f . labelMap

null :: LabelMap a -> Bool
null = IM.null . labelMap

mapMaybe :: (a -> Maybe b) -> LabelMap a -> LabelMap b
mapMaybe f = MkLabelMap . IM.mapMaybe f . labelMap

mapMaybeWithKey :: (Label -> a -> Maybe b) -> LabelMap a -> LabelMap b
mapMaybeWithKey f = MkLabelMap . IM.mapMaybeWithKey (f . keyToLabel) . labelMap

intersection :: LabelMap a -> LabelMap b -> LabelMap a
intersection (MkLabelMap x) (MkLabelMap y) = MkLabelMap (IM.intersection x y)