{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# 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.Proxy
import Data.Some
import Data.Type.Equality
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, TestEquality k) => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjust = adjust' typeCheckK

adjustUnsafe :: Enum1 k => (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjustUnsafe = adjust' don'tCheckK

adjust' :: Enum1 k => Checker k a -> (v a -> v a) -> k a -> DEnumMap k v -> DEnumMap k v
adjust' ch f k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in DEnumMap (IM.adjust (\(KV inf v) -> ch i k inf $ KV inf (f (coe1 v))) i m)

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

alter :: forall k v a. (Enum1 k, TestEquality k) => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alter = alter' typeCheckK

alterUnsafe :: forall k v a. Enum1 k => (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alterUnsafe = alter' don'tCheckK

alter' :: forall k v a. Enum1 k => Checker k a -> (Maybe (v a) -> Maybe (v a)) -> k a -> DEnumMap k v -> DEnumMap k v
alter' ch 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 inf' v)) = ch i k inf' $ KV inf <$> f (Just (coe1 v))

-- alterF

-- * Query
-- ** Lookup

lookup :: (Enum1 k, TestEquality k) => k a -> DEnumMap k v -> Maybe (v a)
lookup = lookup' typeCheckK

lookupUnsafe :: Enum1 k => k a -> DEnumMap k v -> Maybe (v a)
lookupUnsafe = lookup' don'tCheckK

lookup' :: Enum1 k => Checker k a -> k a -> DEnumMap k v -> Maybe (v a)
lookup' ch k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in (\(KV inf v) -> ch i k inf $ coe1 v) <$> IM.lookup i m

-- (!?)
-- (!)

findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a
findWithDefault = findWithDefault' typeCheckK

findWithDefaultUnsafe :: Enum1 k => v a -> k a -> DEnumMap k v -> v a
findWithDefaultUnsafe = findWithDefault' don'tCheckK

findWithDefault' :: Enum1 k => Checker k a -> v a -> k a -> DEnumMap k v -> v a
findWithDefault' ch def k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in case IM.findWithDefault (KV undefined def) i m of
       KV inf' v -> ch i k inf' $ 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 :: DEnumMap k v -> Int
size (DEnumMap m) = IM.size m

-- * Combine

-- ** Union

union :: (Enum1 k, TestEquality k) => DEnumMap k v -> DEnumMap k v -> DEnumMap k v
union = unionWith const  -- if we're checking, we need unionWith anyway, so might as well just delegate here already

-- in the unsafe variant, we can make do with IM.union, which is slightly faster than IM.unionWith, so let's specialise
unionUnsafe :: DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionUnsafe (DEnumMap m1) (DEnumMap m2) = DEnumMap (IM.union m1 m2)

unionWith :: (Enum1 k, TestEquality k)
          => (forall a. v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith f (m1 :: DEnumMap k v) = unionWith' (typeCheckSK (Proxy @k)) f m1

unionWithUnsafe :: (forall a. v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWithUnsafe f (m1 :: DEnumMap k v) = unionWith' (don'tCheckSK (Proxy @k)) f m1

unionWith' :: CheckerSplit k
           -> (forall a. v a -> v a -> v a)
           -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith' ch f (DEnumMap m1 :: DEnumMap k v) (DEnumMap m2) = DEnumMap (IM.unionWithKey f' m1 m2)
  where
    f' :: Int -> KV k v -> KV k v -> KV k v
    f' i (KV inf1 v1) (KV inf2 v2) = ch i inf1 inf2 $ KV inf1 (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 :: Enum1 k => (forall a. k a -> v a -> Bool) -> DEnumMap k v -> (DEnumMap k v, DEnumMap k v)
partitionWithKey f (DEnumMap m) =
  bimap DEnumMap DEnumMap (IM.partitionWithKey (\i (KV inf v) -> case toEnum1 i inf of Some k -> f k (coe1 v)) m)

-- 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


-- * Helpers

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

type CheckerSplit k = forall r. Int -> Enum1Info k -> Enum1Info k -> r -> r

typeCheckSK :: forall k proxy. (Enum1 k, TestEquality k) => proxy k -> CheckerSplit k
typeCheckSK _ i inf1 inf2 = case toEnum1 @k i inf1 of Some k -> typeCheckK i k inf2

don'tCheckSK :: proxy k -> CheckerSplit k
don'tCheckSK _ _ _ _ = id

type Checker k a = forall r. Int -> k a -> Enum1Info k -> r -> r

typeCheckK :: (Enum1 k, TestEquality k) => Checker k a
typeCheckK i k1 inf cont
  | Some k2 <- toEnum1 i inf
  , Just Refl <- testEquality k1 k2 = cont
  | otherwise = errorWithoutStackTrace "DEnumMap: keys with same Int but different types"

don'tCheckK :: Checker k a
don'tCheckK _ _ _ = id