{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Dependent.EnumMap.Strict.Internal where

import Data.Bifunctor (bimap)
import Data.Dependent.Sum
import qualified Data.IntMap.Strict as IM
import Data.Kind (Type)
import Data.Some
import Text.Show (showListWith)
import Unsafe.Coerce (unsafeCoerce)


data KV k v = forall a. KV !(Enum1Info k) !(v a)

newtype DEnumMap k v = DEnumMap (IM.IntMap (KV k v))

instance (Enum1 k, forall a. Show (k a), forall a. Show (v a))
      => Show (DEnumMap (k :: kind -> Type) (v :: kind -> Type)) where
  showsPrec d mp = showParen (d > 10) $
    showString "fromList " . showListWith (\(k :=> v) -> showsPrec 2 k . showString " :=> " . showsPrec 1 v) (toList mp)

class Enum1 f where
  type Enum1Info f
  fromEnum1 :: f a -> (Int, Enum1Info f)
  toEnum1 :: Int -> Enum1Info f -> Some f

-- * Construction

empty :: DEnumMap k v
empty = DEnumMap IM.empty

singleton :: Enum1 k => k a -> v a -> DEnumMap k v
singleton k v =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.singleton i (KV inf v))

-- fromSet

-- ** From Unordered Lists

fromList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromList l =
  DEnumMap (IM.fromList (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))

-- fromListWith
-- fromListWithKey

-- ** From Ascending Lists

-- fromAscList
-- fromAscListWith
-- fromAscListWithKey

fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromDistinctAscList l =
  DEnumMap (IM.fromDistinctAscList (map (\(k :=> v) -> let (i, inf) = fromEnum1 k in (i, KV inf v)) l))

-- * Insertion

insert :: Enum1 k => k a -> v a -> DEnumMap k v -> DEnumMap k v
insert k v (DEnumMap m) =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.insert i (KV inf v) m)

-- insertWith

insertWithKey :: Enum1 k
              => (k a -> v a -> v a -> v a)
              -> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWithKey f k v (DEnumMap m) =
  let (i, inf) = fromEnum1 k
  in DEnumMap (IM.insertWithKey
                (\_ (KV _ v1) (KV _ v2) -> KV inf (f k (coe1 v1) (coe1 v2)))
                i (KV inf v) m)

-- insertLookupWithKey

-- * Deletion\/Update

delete :: Enum1 k => k a -> DEnumMap k v -> DEnumMap k v
delete k (DEnumMap m) = DEnumMap (IM.delete (fst (fromEnum1 k)) m)

adjust :: Enum1 k => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjust f k (DEnumMap m) = DEnumMap (IM.adjust (\(KV k' v) -> KV k' (f (coe1 v))) (fst (fromEnum1 k)) m)

-- adjustWithKey
-- update
-- updateWithKey
-- updateLookupWithKey

alter :: forall k v a. Enum1 k => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alter f k (DEnumMap m) = DEnumMap (IM.alter f' i m)
  where
    (i, inf) = fromEnum1 k

    f' :: Maybe (KV k v) -> Maybe (KV k v)
    f' Nothing = KV inf <$> f Nothing
    f' (Just (KV _ v)) = KV inf <$> f (Just (coe1 v))

-- alterF

-- * Query
-- ** Lookup

lookup :: Enum1 k => k a -> DEnumMap k v -> Maybe (v a)
lookup k (DEnumMap m) = (\(KV _ v) -> coe1 v) <$> IM.lookup (fst (fromEnum1 k)) m

-- (!?)
-- (!)

findWithDefault :: Enum1 k => v a -> k a -> DEnumMap k v -> v a
findWithDefault def k (DEnumMap m) =
  case IM.findWithDefault (KV undefined def) (fst (fromEnum1 k)) m of
    KV _ v -> coe1 v

member :: Enum1 k => k a -> DEnumMap k v -> Bool
member k (DEnumMap m) = IM.member (fst (fromEnum1 k)) m

-- notMember
-- lookupLT
-- lookupGT
-- lookupLE
-- lookupGE

-- ** Size

null :: DEnumMap k v -> Bool
null (DEnumMap m) = IM.null m

-- size

-- * Combine

-- ** Union

union :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
union (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.union m1 m2)

unionWith :: forall k v. (forall a. v a -> v a -> v a)
          -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith f (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.unionWith f' m1 m2)
  where
    f' :: KV k v -> KV k v -> KV k v
    f' (KV inf v1) (KV _ v2) = KV inf (f v1 (coe1 v2))

-- unionWithKey
-- unions
-- unionsWith

-- ** Difference

difference :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
difference (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.difference m1 m2)

-- (\\)
-- differenceWith
-- differenceWithKey

-- ** Intersection

-- intersection
-- intersectionWith
-- intersectionWithKey

-- ** Disjoint

-- disjoint

-- ** Compose

-- compose

-- ** Universal combining function

-- mergeWithKey

-- * Traversal
-- ** Map

-- map
-- mapWithKey
-- traverseWithKey
-- traverseMaybeWithKey
-- mapAccum
-- mapAccumWithKey
-- mapAccumRWithKey
-- mapKeys
-- mapKeysWith
-- mapKeysMonotonic

-- * Folds

-- foldr
-- foldl
-- foldrWithKey
-- foldlWithKey
-- foldMapWithKey

-- ** Strict folds

-- foldr'
-- foldl'
-- foldrWithKey'
-- foldlWithKey'

-- * Conversion

elems :: DEnumMap k v -> [Some v]
elems (DEnumMap m) = map (\(KV _ v) -> Some v) (IM.elems m)

keys :: Enum1 k => DEnumMap k v -> [Some k]
keys (DEnumMap m) = map (\(k, KV inf _) -> toEnum1 k inf) (IM.assocs m)

-- assocs
-- keysSet

-- ** Lists

toList :: Enum1 k => DEnumMap k v -> [DSum k v]
toList = toAscList

-- ** Ordered lists

toAscList :: Enum1 k => DEnumMap k v -> [DSum k v]
toAscList (DEnumMap m) =
  map (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v)
      (IM.toAscList m)

toDescList :: Enum1 k => DEnumMap k v -> [DSum k v]
toDescList (DEnumMap m) =
  map (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v)
      (IM.toDescList m)

-- * Filter

-- filter
-- filterWithKey
-- restrictKeys
-- withoutKeys

partition :: (forall a. v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v) 
partition f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.partition (\(KV _ v) -> f v) m)

-- partitionWithKey

-- takeWhileAntitone
-- dropWhileAntitone
-- spanAntitone

-- mapMaybe
-- mapMaybeWithKey
-- mapEither
-- mapEitherWithKey

-- split
-- splitLookup
-- splitRoot

-- * Submap

-- isSubmapOf, isSubmapOfBy
-- isProperSubmapOf, isProperSubmapOfBy

-- * Min\/Max

-- lookupMin
-- lookupMax
-- findMin
-- findMax
-- deleteMin
-- deleteMax
-- deleteFindMin
-- deleteFindMax
-- updateMin
-- updateMax
-- updateMinWithKey
-- updateMaxWithKey
-- minView
-- maxView
-- minViewWithKey

maxViewWithKey :: Enum1 k => DEnumMap k v -> Maybe (DSum k v, DEnumMap k v)
maxViewWithKey (DEnumMap m) =
  bimap (\(i, KV inf v) -> case toEnum1 i inf of Some k -> k :=> coe1 v) DEnumMap
    <$> IM.maxViewWithKey m


-- * Unsafe helpers

coe1 :: v a -> v b
coe1 = unsafeCoerce