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

import Control.Arrow (second)
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 )
import Control.Applicative ( Applicative )
import Data.Coerce

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

adapt_k_a_m :: Coercible k1 Int =>
                     (Int -> t -> IntMap a -> x) -> k1 -> t -> IMap k2 a -> x
adaptm_k_a_m f k a m = IMap $ adapt_k_a_m f k a m

adaptm_k_a_m :: Coercible k1 Int =>
                      (Int -> t -> IntMap a1 -> IntMap a2)
                      -> k1 -> t -> IMap k2 a1 -> IMap k3 a2
adapt_k_a_m f k a m = adapt_m (adapt_k f k a) m

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)

adaptm_k_m
  :: Coercible k Int =>
     (Int -> IntMap a -> IntMap a) -> k -> IMap k a -> IMap k a
adaptm_k_m f k m = IMap $ adapt_k_m f k 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)

adaptm_m_m :: (IntMap a -> IntMap a -> IntMap a) -> IMap k a -> IMap k a -> IMap k a
adaptm_m_m f a b = IMap $ adapt_m_m f a b

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 :: IMap k a -> Bool
null = adapt_m (IntMap.null)

size :: IMap k a -> Int
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 => k -> IMap k a -> Maybe (k, a)
lookupLT k (IMap m) = coerce $ IntMap.lookupLT (coerce k) m

lookupGT :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
lookupGT k (IMap m) = coerce $ IntMap.lookupGT (coerce k) m

lookupLE :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
lookupLE k (IMap m) = coerce $ IntMap.lookupLE (coerce k) m

lookupGE :: Coercible Int k => k -> IMap k a -> Maybe (k, a)
lookupGE k (IMap m) = coerce $ IntMap.lookupGE (coerce k) m

empty :: IMap k a
empty = IMap IntMap.empty

singleton :: Coercible k Int => k -> a -> IMap k a
singleton = (IMap .) . adapt_k IntMap.singleton

insert :: Coercible k Int => k -> a -> IMap k a -> IMap k a
insert = adaptm_k_a_m IntMap.insert

insertWith :: Coercible k Int => (a -> a -> a) -> k -> a -> IMap k a -> IMap k a
insertWith f = adaptm_k_a_m (IntMap.insertWith f)

insertWithKey :: Coercible Int k => (k -> a -> a -> a) -> k -> a -> IMap k a -> IMap k a
insertWithKey f = adaptm_k_a_m (IntMap.insertWithKey $ f . coerce)

insertLookupWithKey :: Coercible Int k =>
    (k -> a -> a -> a) -> k -> a -> IMap k a -> (Maybe a, IMap k a)
insertLookupWithKey f k a m = second IMap $ adapt_k_a_m (IntMap.insertLookupWithKey $ f . coerce) k a m

delete :: Coercible k Int => k -> IMap k a -> IMap k a
delete = adaptm_k_m IntMap.delete

adjust :: Coercible k Int => (a -> a) -> k -> IMap k a -> IMap k a
adjust f = adaptm_k_m (IntMap.adjust f)

adjustWithKey :: Coercible Int k =>
     (k -> a -> a) -> k -> IMap k a -> IMap k a
adjustWithKey f = adaptm_k_m (IntMap.adjustWithKey $ f . coerce)

update
  :: Coercible k Int => (a -> Maybe a) -> k -> IMap k a -> IMap k a
update f = adaptm_k_m (IntMap.update f)


updateWithKey :: Coercible Int k =>
     (k -> a -> Maybe a) -> k -> IMap k a -> IMap k a
updateWithKey f = adaptm_k_m (IntMap.updateWithKey $ f . coerce)

updateLookupWithKey :: Coercible k Int =>
  (k -> a -> Maybe a) -> k -> IMap k a -> (Maybe a, IMap k a)
updateLookupWithKey f k m =
    second IMap $ adapt_k_m (IntMap.updateLookupWithKey $ f . coerce) k m

alter :: Coercible k Int => (Maybe a -> Maybe a) -> k -> IMap k a -> IMap k a
alter f = adaptm_k_m (IntMap.alter f)

union :: IMap k a -> IMap k a -> IMap k a
union = adaptm_m_m IntMap.union

unionWith :: (a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
unionWith f = adaptm_m_m (IntMap.unionWith f)


unionWithKey :: Coercible Int k => (k -> a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
unionWithKey f = adaptm_m_m (IntMap.unionWithKey $ f . coerce)

unions :: Coercible k Int => [IMap k a] -> IMap k a
unions ms = IMap $ IntMap.unions (coerce <$> ms)

unionsWith :: Coercible k Int => (a->a->a) -> [IMap k a] -> IMap k a
unionsWith f ms = IMap $ IntMap.unionsWith f (coerce <$> ms)

difference :: IMap k b -> IMap k b -> IMap k b
difference = adaptm_m_m IntMap.difference

differenceWith :: (b -> b -> Maybe b)
                        -> IMap k b -> IMap k b -> IMap k b
differenceWith f = adaptm_m_m (IntMap.differenceWith f)

differenceWithKey ::
  Coercible Int k =>
  (k -> a -> a -> Maybe a) -> IMap k a -> IMap k a -> IMap k a
differenceWithKey f = adaptm_m_m (IntMap.differenceWithKey $ f . coerce)

intersection :: IMap k b -> IMap k b -> IMap k b
intersection = adaptm_m_m IntMap.intersection

intersectionWith :: (a -> a -> a) -> IMap k a -> IMap k a -> IMap k a
intersectionWith f = adaptm_m_m (IntMap.intersectionWith f)

mergeWithKey ::
  Coercible Int k =>
  (k -> a -> b -> Maybe c)
  -> (IMap k a -> IMap k c)
  -> (IMap k b -> IMap k c)
  -> IMap k a
  -> IMap k b
  -> IMap k c
mergeWithKey f g1 g2 = adaptm_m_m (IntMap.mergeWithKey f' g1' g2')
 where f'  = f . coerce
       g1' = intmap . g1 . IMap
       g2' = intmap . g2 . IMap
       adapt_m_m f m = adapt_m (adapt_m f m)
       adaptm_m_m f a b = IMap $ adapt_m_m f a b

map :: (a -> b) -> IMap k a -> IMap k b
map f = IMap . adapt_m (IntMap.map f)

mapWithKey :: Coercible Int k => (k -> a -> b) -> IMap k a -> IMap k b
mapWithKey f = IMap . adapt_m (IntMap.mapWithKey $ f . coerce)

-- FIXME: fmap IMap ?
traverseWithKey ::
  (Applicative f, Coercible Int k) =>
  (k -> a -> f b) -> IMap k a -> f (IMap k b)
traverseWithKey f = fmap IMap . adapt_m (IntMap.traverseWithKey $ f . coerce)

mapAccum :: (t -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
mapAccum f a m = second IMap $ IntMap.mapAccum f a (intmap m)

mapAccumWithKey :: Coercible Int k =>
    (t -> k -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
mapAccumWithKey f a m = second IMap $ IntMap.mapAccumWithKey f' a (intmap m)
 where f' a k b = f a (coerce k) b


mapAccumRWithKey :: Coercible Int k =>
    (t -> k -> b -> (t, a)) -> t -> IMap k b -> (t, IMap k a)
mapAccumRWithKey f a m = second IMap $ IntMap.mapAccumRWithKey f' a (intmap m)
 where f' a k b = f a (coerce k) b

mapKeys :: (Coercible Int k1, Coercible Int k2) =>
    (k1 -> k2) -> IMap k1 a -> IMap k2 a
mapKeys f = IMap . adapt_m (IntMap.mapKeys (coerce . f . coerce))

mapKeysWith :: (Coercible Int k1, Coercible Int k2) =>
    (a->a->a) -> (k1 -> k2) -> IMap k1 a -> IMap k2 a
mapKeysWith c f = IMap . adapt_m (IntMap.mapKeysWith c (coerce . f . coerce))

mapKeysMonotonic :: (Coercible Int k1, Coercible Int k2) =>
    (k1 -> k2) -> IMap k1 a -> IMap k2 a
mapKeysMonotonic f = IMap . adapt_m (IntMap.mapKeysMonotonic (coerce . f . coerce))

foldr :: (a -> x -> x) -> x -> IMap k a -> x
foldr f b = adapt_m (IntMap.foldr f b)

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

foldrWithKey :: Coercible Int b => (b -> a -> x -> x) -> x -> IMap k a -> x
foldrWithKey f b = adapt_m (IntMap.foldrWithKey (f . coerce) b)

foldlWithKey ::
  Coercible Int k => (x -> k -> a -> x) -> x -> IMap k a -> x
foldlWithKey f a = adapt_m (IntMap.foldlWithKey f' a) where f' a = f a . coerce

foldMapWithKey :: (Monoid m, Coercible Int k) => (k -> a -> m) -> IMap k a -> m
foldMapWithKey f = adapt_m (IntMap.foldMapWithKey $ f . coerce)

foldr' :: (a -> x -> x) -> x -> IMap k a -> x
foldr' f b = adapt_m (IntMap.foldr' f b)

foldl' :: (a -> x -> a) -> a -> IMap k x -> a
foldl' f b = adapt_m (IntMap.foldl' f b)

foldrWithKey' :: Coercible Int b => (b -> a -> x -> x) -> x -> IMap k a -> x
foldrWithKey' f b = adapt_m (IntMap.foldrWithKey' (f . coerce) b)

foldlWithKey' ::
  Coercible Int k => (x -> k -> a -> x) -> x -> IMap k a -> x
foldlWithKey' f a = adapt_m (IntMap.foldlWithKey' f' a) where f' a = f a . coerce

elems :: IMap k a -> [a]
elems = IntMap.elems . intmap

keys :: Coercible Int k => IMap k a -> [k]
keys = coerce . IntMap.keys . intmap

assocs :: Coercible Int k => IMap k a -> [(k, a)]
assocs = coerce . IntMap.assocs . intmap

-- Not implementing... (doing it right requires wrapping IntSet)
-- keysSet :: IntMap a -> IntSet
-- fromSet :: (Key -> a) -> IntSet -> IntMap a

toList :: Coercible Int k => IMap k a -> [(k, a)]
toList = coerce . IntMap.toList . intmap

fromList :: Coercible Int k => [(k, a)] -> IMap k a
fromList = IMap . IntMap.fromList . coerce

fromListWith :: Coercible Int k => (a -> a -> a) -> [(k, a)] -> IMap k a
fromListWith f = IMap . IntMap.fromListWith f . coerce

fromListWithKey :: Coercible Int k =>
    (k -> a -> a -> a) -> [(k, a)] -> IMap k a
fromListWithKey f = IMap . IntMap.fromListWithKey (f . coerce) . coerce

toAscList :: Coercible Int k => IMap k a -> [(k,a)]
toAscList (IMap m) = coerce $ IntMap.toAscList m

toDescList :: Coercible Int k => IMap k a -> [(k,a)]
toDescList (IMap m) = coerce $ IntMap.toDescList m

fromAscList :: Coercible Int k => [(k, a)] -> IMap k a
fromAscList = IMap . IntMap.fromAscList . coerce

fromAscListWith :: Coercible Int k
  => (a -> a -> a) -> [(k,a)] -> IMap k a
fromAscListWith f = IMap . IntMap.fromAscListWith f . coerce

fromAscListWithKey :: Coercible Int k
  => (k -> a -> a -> a) -> [(k,a)] -> IMap k a
fromAscListWithKey f = IMap . IntMap.fromAscListWithKey (f . coerce) . coerce

fromDistinctAscList :: Coercible Int k => [(k, a)] -> IMap k a
fromDistinctAscList = IMap . IntMap.fromDistinctAscList . coerce

filter :: (a -> Bool) -> IMap k a -> IMap k a
filter f = IMap . adapt_m (IntMap.filter f)

filterWithKey :: Coercible Int k => (k -> a -> Bool) -> IMap k a -> IMap k a
filterWithKey f = IMap . adapt_m (IntMap.filterWithKey $ f . coerce)

partition :: (a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
partition f m = coerce $ IntMap.partition f (intmap m)


partitionWithKey :: Coercible Int k
    => (k -> a -> Bool) -> IMap k a -> (IMap k a, IMap k a)
partitionWithKey f m = coerce $ IntMap.partitionWithKey (f . coerce) (intmap m)

mapMaybe :: (a -> Maybe b) -> IMap k a -> IMap k b
mapMaybe f = IMap . IntMap.mapMaybe f . intmap