{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveDataTypeable        #-}
{-# LANGUAGE DeriveFoldable            #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ExplicitNamespaces        #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TupleSections             #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UndecidableInstances      #-}


module Data.Set.Of
  ( SetOf
  , toSet
  , fromFoldable
  , fromFoldable'
  , Data.Set.Of.toList
  , elems
  , SetToTuple(..)
  , member
  , notMember
  , fst
  , snd
  , findMin
  , findMax
#if MIN_VERSION_containers(0,5,9)
  , lookupMax
  , lookupMin
#endif
  , map
  , mapMonotonic
  , fold
  , foldl'
  , foldr'
  , module GHC.TypeLits
  , empty
  , singleton
  , doubleton
  , tripleton
  , uncheckedmkSetOf
  , SomeSetOf (..)
  , union
  , withSomeSetOf
  ) where


import           Data.Data
import qualified Data.Foldable as F
import           Data.Kind
import           Data.List     (sort)
import           Data.Monoid   ((<>))
import           Data.Set      (Set)
import qualified Data.Set      as S
import           GHC.Exts
import           GHC.Generics
import           GHC.TypeLits  (type (+), KnownNat, Nat, SomeNat (..), natVal,
                                someNatVal)
import           Prelude       hiding (fst, map, snd)


-- | A set of values @a@ with a size of @n@
newtype SetOf (n :: Nat) a = SetOf { toSet :: Set a }
  deriving (Show, Foldable, Eq, Ord, Data, Generic)


instance (Ord a, KnownNat n) => IsList (Maybe (SetOf n a)) where
  type Item (Maybe (SetOf n a)) = a
  fromList = fromFoldable @ n
  toList = maybe [] Data.Set.Of.toList


-- | /O(n*log n)/. Create a set from a list of elements.
fromFoldable' :: forall f a (n :: Nat). (KnownNat n, Ord a, Foldable f) => Proxy n -> f a -> Maybe (SetOf n a)
fromFoldable' n xs = let sxs = S.fromList $ F.toList xs in
  if F.length sxs == fromIntegral (natVal n)
  then Just $ SetOf sxs else Nothing


fromFoldable :: forall (n :: Nat) f a. (KnownNat n, Ord a, Foldable f) => f a -> Maybe (SetOf n a)
fromFoldable = fromFoldable' (Proxy @ n)


-- | /O(n)/. Convert the set to a list of elements. Subject to list fusion.
toList, elems :: Ord a => SetOf n a -> [a]
toList (SetOf s) = sort $ S.toList s
elems = Data.Set.Of.toList


-- | The empty set
empty :: SetOf 0 a
empty = SetOf S.empty


-- |O(1). Create a singleton set.
singleton :: a -> SetOf 1 a
singleton = SetOf . S.singleton


doubleton :: Ord a => a -> a -> Maybe (SetOf 2 a)
doubleton = curry fromTuple


tripleton :: Ord a => a -> a -> a -> Maybe (SetOf 3 a)
tripleton x y z = fromTuple (x, y, z)


-- | Type class to provide type safe, polymorphic mappings to and from tuples
class SetToTuple (n :: Nat) a where
  -- | The corrosponding tuple, as a family
  type TupOf n a :: Type
  -- | Take a tuple, and potentially produce a Set of a the given size.
  -- @Nothing@ is produced when the input tuple contains duplicate elements.
  fromTuple :: Ord a => TupOf n a -> Maybe (SetOf n a)
  -- | Produce a tuple, where the elements @a@ are ordered left to right descending.
  toTuple :: Ord a => SetOf n a -> TupOf n a


instance SetToTuple 0 a where
  -- | We treat @()@ as a zerotuple
  type TupOf 0 a = ()
  fromTuple () = fromFoldable @ 0 []
  toTuple s = case Data.Set.Of.toList s of
              [] -> ()
              _  -> error "Not a valid SetOf 0"


instance SetToTuple 1 a where
  type TupOf 1 a = a
  fromTuple x = fromFoldable @ 1 [x]
  toTuple s = case Data.Set.Of.toList s of
              [x] -> x
              _   -> error "Not a valid SetOf 1"


instance SetToTuple 2 a where
  type TupOf 2 a = (a, a)
  fromTuple (x, y) = fromFoldable @ 2 [x, y]
  toTuple s = case Data.Set.Of.toList s of
              [x, y] -> (x, y)
              _      -> error "Not a valid SetOf 2"


instance SetToTuple 3 a where
  type TupOf 3 a = (a, a, a)
  fromTuple (x, y, z) = fromFoldable @ 3 [x, y, z]
  toTuple s = case Data.Set.Of.toList s of
              [x, y, z] -> (x, y, z)
              _         -> error "Not a valid SetOf 3"


instance SetToTuple 4 a where
  type TupOf 4 a = (a, a, a, a)
  fromTuple (w, x, y, z) = fromFoldable @ 4 [w, x, y, z]
  toTuple s = case Data.Set.Of.toList s of
              [w, x, y, z] -> (w, x, y, z)
              _            -> error "Not a valid SetOf 4"


instance SetToTuple 5 a where
  type TupOf 5 a = (a, a, a, a, a)
  fromTuple (v, w, x, y, z) = fromFoldable @ 5 [v, w, x, y, z]
  toTuple s = case Data.Set.Of.toList s of
              [v, w, x, y, z] -> (v, w, x, y, z)
              _               -> error "Not a valid SetOf 5"


instance SetToTuple 6 a where
  type TupOf 6 a = (a, a, a, a, a, a)
  fromTuple (u, v, w, x, y, z) = fromFoldable @ 6 [u, v, w, x, y, z]
  toTuple s = case Data.Set.Of.toList s of
              [u, v, w, x, y, z] -> (u, v, w, x, y, z)
              _                  -> error "Not a valid SetOf 6"


instance SetToTuple 7 a where
  type TupOf 7 a = (a, a, a, a, a, a, a)
  fromTuple (t, u, v, w, x, y, z) = fromFoldable @ 7 [t, u, v, w, x, y, z]
  toTuple s = case Data.Set.Of.toList s of
              [t, u, v, w, x, y, z] -> (t, u, v, w, x, y, z)
              _                     -> error "Not a valid SetOf 7"


-- | /O(log n)/. Is the element in the set?
member :: Ord a => a -> SetOf n a -> Bool
member x = S.member x . toSet


-- | /O(log n)/. Is the element not in the set?
notMember :: Ord a => a -> SetOf n a -> Bool
notMember x = S.notMember x . toSet


-- | /O(log n)/. The maximal element of a set.
fst :: Ord a => SetOf 2 a -> a
fst = (\(x,_)->x) . toTuple


-- | /O(log n)/. The minimal element of a set.
snd :: Ord a => SetOf 2 a -> a
snd = (\(_,x)->x) . toTuple


#if MIN_VERSION_containers(0,5,9)
lookupMax, lookupMin :: SetOf n a -> Maybe a
lookupMax = S.lookupMax . toSet
lookupMin = S.lookupMin . toSet
#endif


-- | Now typesafe!
findMax, findMin :: SetOf (n + 1) a -> a
findMax = S.findMax . toSet
findMin = S.findMin . toSet


-- | /O(n*log n)/.
-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@.
--
-- It's worth noting that the size of the result may be smaller if,
-- for some @(x,y)@, @x \/= y && f x == f y@
--
-- If the size changed, you will get @Nothing@.
map :: forall n a b. (KnownNat n, Ord b) => (a -> b) -> SetOf n a -> Maybe (SetOf n b)
map f = fromFoldable @ n . S.map f . toSet


-- | /O(n)/. The
--
-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly increasing.
-- /The precondition is not checked./
-- Semi-formally, we have:
--
-- > and [x < y ==> f x < f y | x <- ls, y <- ls]
-- >                     ==> mapMonotonic f s == map f s
-- >     where ls = toList s
mapMonotonic :: (a -> b) -> SetOf n a -> SetOf n b
mapMonotonic f (SetOf s) = SetOf $ S.mapMonotonic f s


-- | /O(n)/. A strict version of 'foldl'. Each application of the operator is
-- evaluated before using the result in the next application. This
-- function is strict in the starting value.<Paste>
foldl' :: (a -> b -> a) -> a -> SetOf n b -> a
foldl' f x = S.foldl' f x . toSet
{-# INLINE foldl' #-}


-- | /O(n)/. A strict version of 'foldr'. Each application of the operator is
-- evaluated before using the result in the next application. This
-- function is strict in the starting value.
foldr' :: (a -> b -> b) -> b -> SetOf n a -> b
foldr' f x = S.foldr' f x . toSet
{-# INLINE foldr' #-}


-- | /O(n)/. Fold the elements in the set using the given right-associative
-- binary operator. This function is an equivalent of 'foldr' and is present
-- for compatibility only.
--
-- /Please note that fold will be deprecated in the future and removed./
fold :: (a -> b -> b) -> b -> SetOf n a -> b
fold f x = S.fold f x . toSet
{-# INLINE fold #-}



uncheckedmkSetOf :: Foldable f => f a -> SetOf n a
uncheckedmkSetOf = error "boo, you thought you could cheat SetOf, but you can't."


data SomeSetOf a = forall n. SomeSetOf (SetOf n a)


instance Show a => Show (SomeSetOf a) where
  show (SomeSetOf s) = "SomeSetOf (" <> show s <> ")"


instance Eq a => Eq (SomeSetOf a) where
  SomeSetOf (SetOf s) == SomeSetOf (SetOf s') = s == s'


instance Ord a => Ord (SomeSetOf a) where
  compare (SomeSetOf (SetOf s)) (SomeSetOf (SetOf s')) = compare s s'





union :: Ord a => SetOf n a -> SetOf n' a -> SomeSetOf a
union (SetOf x) (SetOf y) = SomeSetOf . SetOf $ x <> y


withSomeSetOf :: forall a r. (forall n. Proxy n -> SetOf n a -> r) -> SomeSetOf a -> Maybe r
withSomeSetOf f (SomeSetOf (SetOf s)) =
  (\(SomeNat n) -> f n (SetOf s)) <$> someNatVal (fromIntegral $ S.size s)