{-# LANGUAGE CPP,
             FlexibleContexts,
             MultiParamTypeClasses,
             GeneralizedNewtypeDeriving,
             DeriveTraversable,
             DeriveDataTypeable #-}
module IntMapClass where

import qualified Data.IntMap.Strict as IntMap
import Data.IntMap.Strict ( IntMap )
import Data.Typeable ( Typeable )
import Data.Data     ( Data )
import Data.Foldable ( Foldable )
import Data.Traversable ( Traversable )
import Data.Monoid   ( Monoid )
import Control.DeepSeq ( NFData )
#if MIN_VERSION_base(4,7,0)
import Data.Coerce
#else
class Coercible a b where coerce :: a -> b
#endif

newtype IMap k a = IMap { intmap :: IntMap a }
 deriving
    ( Functor
    , Typeable
    , Foldable
    , Traversable
    , Eq
    , Data
    , Ord
    , Read
    , Show
    , Monoid
    , NFData
    )

adapt_m_k :: Coercible k Int => (IntMap a -> Int -> x) -> IMap k a -> k -> x
adapt_m_k f (IMap m) k = f m (coerce k)

adapt_k_m :: Coercible k Int => (Int -> IntMap a -> x) -> k -> IMap k a -> x
adapt_k_m f k (IMap m) = f (coerce k) m
-- adapt_k_m2 :: Coercible k Int => (Int -> IntMap a -> x) -> k -> IMap k a -> x
-- adapt_k_m2 f k m = (adapt_k f) k (intmap m)
-- adapt_k :: Coercible k Int => (Int -> x) -> k -> x
-- adapt_k f k = f (coerce k)

adapt_m_m :: (IntMap a -> IntMap a -> x) -> IMap k a -> IMap k a -> x
adapt_m_m f m = adapt_m (adapt_m f m)

adapt_m :: (IntMap a -> x) -> IMap k a -> x
adapt_m f (IMap m) = f m


(!) :: Coercible k Int => IMap k a -> k -> a
(!) = adapt_m_k (IntMap.!)

(\\) :: IMap k a -> IMap k a -> IMap k a
(\\) a b = IMap $ adapt_m_m (IntMap.\\) a b

null = adapt_m (IntMap.null)
size = adapt_m (IntMap.size)

member :: Coercible k Int => k -> IMap k a -> Bool
member = adapt_k_m (IntMap.member)

notMember :: Coercible k Int => k -> IMap k a -> Bool
notMember = adapt_k_m (IntMap.notMember)

lookup :: Coercible k Int => k -> IMap k a -> Maybe a
lookup = adapt_k_m (IntMap.lookup)

findWithDefault :: Coercible k Int => x -> k -> IMap k x -> x
findWithDefault a = adapt_k_m (IntMap.findWithDefault a)

lookupLT :: ( Coercible Int k, Coercible k Int ) => k -> IMap k a -> Maybe (k, a)
lookupLT k m = fmap (first coerce) $ adapt_k_m (IntMap.lookupLT) k m
               where first f (x,y) = (f x,y)