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

import Control.Exception
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)

-- | This class attempts to generalise 'Enum' to indexed data types: data types
-- with a GADT-like type parameter. Conversion to an 'Int' naturally loses type
-- information, and furthermore it is common to actually need some additional
-- data alongside the 'Int' to be able to reconstruct the original (in
-- 'toEnum1'). This additional data lives in 'Enum1Info'. The laws are:
--
-- [Unique IDs]
--   If @'fst' ('fromEnum1' x) == 'fst' ('fromEnum1' y)@ then @'testEquality' x y == 'Just' 'Refl' && x '==' y@
-- [Persistent IDs]
--   @'fst' ('fromEnum1' ('uncurry' 'toEnum1' ('fromEnum1' x))) == 'fst' ('fromEnum1' x)@
--
-- The "Unique IDs" law states that if the IDs of two values are equal, then
-- the values themselves must have the same type index, and furthermore be
-- equal. If @f@ does not implement 'TestEquality' or 'Eq', the law should
-- morally hold (but most of the API will be unusable).
--
-- The "Persistent IDs" law states that reconstructing a value using 'toEnum1'
-- does not change its ID.
--
-- __Note__: The methods on 'DEnumMap' attempt to check these laws using
-- 'assert' assertions (which are by default __disabled__ when optimisations
-- are on!), but full consistency cannot always be checked; if you break these
-- laws in a sufficiently clever way, the internals of 'DEnumMap' may
-- 'unsafeCoerce' unequal things and engage nasal demons, including crashes and
-- worse.
class Enum1 f where
  type Enum1Info f
  fromEnum1 :: f a -> (Int, Enum1Info f)
  toEnum1 :: Int -> Enum1Info f -> Some f

dSumToKV :: Enum1 k => DSum k v -> (Int, KV k v)
dSumToKV (k :=> v) = let (i, inf) = fromEnum1 k in (i, KV inf v)

-- | Assumes that the input was obtained via 'fromEnum1'.
kVToDSum :: Enum1 k => (Int, KV k v) -> DSum k v
kVToDSum (i, KV inf v) = case toEnum1 i inf of Some k -> k :=> coe1 v

-- * 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 dSumToKV l))

fromListWith :: (Enum1 k, TestEquality k)
             => (forall a. v a -> v a -> v a)
             -> [DSum k v] -> DEnumMap k v
fromListWith f (l :: [DSum k v]) =
  DEnumMap (IM.fromListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
                typeCheck2 (Proxy @k) i inf1 inf2 $
                  KV inf1 (f v1 (coe1 v2)))
             (map dSumToKV l))

fromListWithKey :: (Enum1 k, TestEquality k)
                => (forall a. k a -> v a -> v a -> v a)
                -> [DSum k v] -> DEnumMap k v
fromListWithKey f l =
  DEnumMap (IM.fromListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
                case toEnum1 i inf1 of
                  Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2)))
             (map dSumToKV l))

-- ** From Ascending Lists

fromAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromAscList l = DEnumMap (IM.fromAscList (map dSumToKV l))

fromAscListWith :: (Enum1 k, TestEquality k)
                => (forall a. v a -> v a -> v a)
                -> [DSum k v] -> DEnumMap k v
fromAscListWith f (l :: [DSum k v]) =
  DEnumMap (IM.fromAscListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
               typeCheck2 (Proxy @k) i inf1 inf2 $
                 KV inf1 (f v1 (coe1 v2)))
             (map dSumToKV l))

fromAscListWithKey :: (Enum1 k, TestEquality k)
                   => (forall a. k a -> v a -> v a -> v a)
                   -> [DSum k v] -> DEnumMap k v
fromAscListWithKey f l =
  DEnumMap (IM.fromAscListWithKey
             (\i (KV inf1 v1) (KV inf2 v2) ->
               case toEnum1 i inf1 of
                 Some k1 -> typeCheck1 k1 i inf2 $ KV inf1 (f k1 (coe1 v1) (coe1 v2)))
             (map dSumToKV l))

fromDistinctAscList :: Enum1 k => [DSum k v] -> DEnumMap k v
fromDistinctAscList l = DEnumMap (IM.fromDistinctAscList (map dSumToKV 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 :: (Enum1 k, TestEquality k)
           => (v a -> v a -> v a)
           -> k a -> v a -> DEnumMap k v -> DEnumMap k v
insertWith = insertWithKey . const

insertWithKey :: (Enum1 k, TestEquality 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.insertWith
                (\_ (KV inf' v2) -> typeCheck1 k i inf' $ KV inf (f k v (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 f k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in DEnumMap (IM.adjust (\(KV inf v) -> typeCheck1 k i 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 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)) = typeCheck1 k i inf' $ KV inf <$> f (Just (coe1 v))

-- alterF

-- * Query
-- ** Lookup

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

-- (!?)
-- (!)

findWithDefault :: (Enum1 k, TestEquality k) => v a -> k a -> DEnumMap k v -> v a
findWithDefault def k (DEnumMap m) =
  let (i, _) = fromEnum1 k
  in case IM.findWithDefault (KV undefined def) i m of
       KV inf' v -> typeCheck1 k i 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 type checking, we need unionWith anyway, so might as well just delegate here already

unionWith :: (Enum1 k, TestEquality k)
          => (forall a. v a -> v a -> v a) -> DEnumMap k v -> DEnumMap k v -> DEnumMap k v
unionWith 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) = typeCheck2 (Proxy @k) 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 kVToDSum (IM.toAscList m)

toDescList :: Enum1 k => DEnumMap k v -> [DSum k v]
toDescList (DEnumMap m) = map kVToDSum (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 kVToDSum DEnumMap <$> IM.maxViewWithKey m


-- * Helpers

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

typeCheck1 :: (Enum1 k, TestEquality k)
           => k a -> Int -> Enum1Info k -> r -> r
typeCheck1 k1 i inf2 x =
  assert (case toEnum1 i inf2 of { Some k2 ->
          case testEquality k1 k2 of
            Just Refl -> True
            Nothing -> False })
         x

typeCheck2 :: forall k proxy r. (Enum1 k, TestEquality k)
           => proxy k -> Int -> Enum1Info k -> Enum1Info k -> r -> r
typeCheck2 _ i inf1 inf2 x =
  assert (case toEnum1 @k i inf1 of { Some k1 ->
          case toEnum1 i inf2 of { Some k2 ->
          case testEquality k1 k2 of
            Just Refl -> True
            Nothing -> False }})
         x