--------------------------------------------------------------------------------
-- |
-- Module      :  HGeometry.Set.Util
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module HGeometry.Set.Util
  ( toggle

  , splitOn, splitBy
  , fromListBy
  , join
  , insertBy
  , deleteAllBy
  , queryBy
  , toggleBy
  ) where

import           Data.Functor.Identity
import           Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Set.Internal as Internal
import           HGeometry.Ord.Dynamic

--------------------------------------------------------------------------------

-- | Insert the element if it is present, or delete it otherwise.
--
-- >>> toggle 5 $ Set.fromList [1,2,3,6]
-- fromList [1,2,3,5,6]
-- >>> toggle 5 $ Set.fromList [1,2,3,5,6]
-- fromList [1,2,3,6]
toggle   :: Ord a => a -> Set a -> Set a
toggle :: forall a. Ord a => a -> Set a -> Set a
toggle a
x = Identity (Set a) -> Set a
forall a. Identity a -> a
runIdentity (Identity (Set a) -> Set a)
-> (Set a -> Identity (Set a)) -> Set a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Identity Bool) -> a -> Set a -> Identity (Set a)
forall a (f :: * -> *).
(Ord a, Functor f) =>
(Bool -> f Bool) -> a -> Set a -> f (Set a)
Set.alterF (Bool -> Identity Bool
forall a. a -> Identity a
Identity (Bool -> Identity Bool) -> (Bool -> Bool) -> Bool -> Identity Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not) a
x


--------------------------------------------------------------------------------
-- * Operations related to consistent orderings

-- data S = S String deriving Show
-- cmpS :: S -> S -> Ordering
-- cmpS = comparing (\(S s) -> length s)


-- $setup
-- >>> import Data.Ord(comparing)
-- >>> data S = S String deriving Show
-- >>> cmpS = comparing (\(S s) -> length s)
--

-- | Given a monotonic function f that maps a to b, and a value x, split the sequence s
-- depending on the b values. I.e. the result (l,m,r) is such that
-- * all (\y -> f y <  x) l
-- * all (\y -> f y == x) m
-- * all (\y -> f y >  x) r
--
-- running time: \(O(\log n)\)
splitOn       :: Ord b => (a -> b) -> b -> Set a -> (Set a, Set a, Set a)
splitOn :: forall b a.
Ord b =>
(a -> b) -> b -> Set a -> (Set a, Set a, Set a)
splitOn a -> b
f b
x Set a
s = let (Set a
l,Set a
s') = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> b -> Bool
g Ordering
LT (b -> Bool) -> (a -> b) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) Set a
s
                    (Set a
m,Set a
r)  = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> b -> Bool
g Ordering
EQ (b -> Bool) -> (a -> b) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f) Set a
s'
                    g :: Ordering -> b -> Bool
g Ordering
c b
y  = b
y b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` b
x Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
c
                in (Set a
l,Set a
m,Set a
r)

-- | Given a monotonic function f that orders @a@, split the sequence @s@
-- into three parts. I.e. the result (lt,eq,gt) is such that
-- * all (\x -> f x == LT) lt
-- * all (\x -> f x == EQ) eq
-- * all (\x -> f x == GT) gt
--
-- running time: \(O(\log n)\)
splitBy       :: (a -> Ordering) -> Set a -> (Set a, Set a, Set a)
splitBy :: forall a. (a -> Ordering) -> Set a -> (Set a, Set a, Set a)
splitBy a -> Ordering
f Set a
s = let (Set a
l,Set a
s') = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
(==) Ordering
LT (Ordering -> Bool) -> (a -> Ordering) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Ordering
f) Set a
s
                  (Set a
m,Set a
r)  = (a -> Bool) -> Set a -> (Set a, Set a)
forall a. (a -> Bool) -> Set a -> (Set a, Set a)
Set.spanAntitone (Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
(==) Ordering
EQ (Ordering -> Bool) -> (a -> Ordering) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Ordering
f) Set a
s'
              in (Set a
l,Set a
m,Set a
r)

-- | Constructs a Set using the given Order.
--
-- Note that this is dangerous as the resulting set may not abide the
-- ordering expected of such sets.
--
-- running time: \(O(n\log n)\)
fromListBy        :: (a -> a -> Ordering) -> [a] -> Set a
fromListBy :: forall a. (a -> a -> Ordering) -> [a] -> Set a
fromListBy a -> a -> Ordering
cmp [a]
xs = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp (Set (O s a) -> O s (Set a)
forall (f :: * -> *) s a. f (O s a) -> O s (f a)
extractOrd1 (Set (O s a) -> O s (Set a))
-> ([a] -> Set (O s a)) -> [a] -> O s (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [O s a] -> Set (O s a)
forall a. Ord a => [a] -> Set a
Set.fromList ([O s a] -> Set (O s a)) -> ([a] -> [O s a]) -> [a] -> Set (O s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> O s a) -> [a] -> [O s a]
forall a b. (a -> b) -> [a] -> [b]
map a -> O s a
forall s a. a -> O s a
O ([a] -> O s (Set a)) -> [a] -> O s (Set a)
forall a b. (a -> b) -> a -> b
$ [a]
xs)

-- | Given two sets l and r, such that all elements of l occur before
-- r, join the two sets into a combined set.
--
-- running time: \(O(\log n)\)
join :: Set a -> Set a -> Set a
join :: forall a. Set a -> Set a -> Set a
join = Set a -> Set a -> Set a
forall a. Set a -> Set a -> Set a
Internal.merge


-- | Inserts an element into the set, assuming that the set is ordered
-- by the given order.
--
-- >>> insertBy cmpS (S "ccc") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "bb",S "ccc",S "dddd"]
--
-- When trying to insert an element that equals an element already in
-- the set (according to the given comparator), this function replaces
-- the old element by the new one:
--
-- >>> insertBy cmpS (S "cc") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "cc",S "dddd"]
--
-- running time: \(O(\log n)\)
insertBy         :: (a -> a -> Ordering) -> a -> Set a -> Set a
insertBy :: forall a. (a -> a -> Ordering) -> a -> Set a -> Set a
insertBy a -> a -> Ordering
cmp a
x Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> Set (O s a)) -> Set a -> O s (Set a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> Set (O s a)
forall a. Ord a => a -> Set a -> Set a
Set.insert (O s a -> Set (O s a) -> Set (O s a))
-> O s a -> Set (O s a) -> Set (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
x) Set a
s


-- | Deletes an element from the set, assuming the set is ordered by
-- the given ordering.
--
-- >>> deleteAllBy cmpS (S "bb") $ fromListBy cmpS [S "a" , S "bb" , S "dddd"]
-- fromList [S "a",S "dddd"]
-- >>> deleteAllBy cmpS (S "bb") $ fromListBy cmpS [S "a" , S "bb" , S "cc", S "dd", S "ee", S "ff", S "dddd"]
-- fromList [S "a",S "dddd"]
--
-- running time: \(O(\log n)\)
deleteAllBy         :: (a -> a -> Ordering) -> a -> Set a -> Set a
deleteAllBy :: forall a. (a -> a -> Ordering) -> a -> Set a -> Set a
deleteAllBy a -> a -> Ordering
cmp a
x Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> Set (O s a)) -> Set a -> O s (Set a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> Set (O s a)
forall a. Ord a => a -> Set a -> Set a
Set.delete (O s a -> Set (O s a) -> Set (O s a))
-> O s a -> Set (O s a) -> Set (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
x) Set a
s

-- | Flips an element by the given order
toggleBy         :: (a -> a -> Ordering) -> a -> Set a -> Set a
toggleBy :: forall a. (a -> a -> Ordering) -> a -> Set a -> Set a
toggleBy a -> a -> Ordering
cmp a
x Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a)
-> (forall s. Reifies s (OrdDict a) => O s (Set a)) -> Set a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> Set (O s a)) -> Set a -> O s (Set a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> Set (O s a)
forall a. Ord a => a -> Set a -> Set a
toggle (O s a -> Set (O s a) -> Set (O s a))
-> O s a -> Set (O s a) -> Set (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
x) Set a
s

-- | Run a query, eg. lookupGE, on the set with the given ordering.
--
-- Note: The 'Algorithms.BinarySearch.binarySearchIn' function may be
-- a useful alternative to 'queryBy'
--
-- >>> queryBy cmpS Set.lookupGE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "bbb")
-- >>> queryBy cmpS Set.lookupLE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "a")
-- >>> queryBy cmpS Set.lookupGE (S "333") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- Just (S "bbb")
queryBy           :: (a -> a -> Ordering)
                  -> (forall b. Ord b => b -> Set b -> t b)
                  -> a -> Set a -> t a
queryBy :: forall a (t :: * -> *).
(a -> a -> Ordering)
-> (forall b. Ord b => b -> Set b -> t b) -> a -> Set a -> t a
queryBy a -> a -> Ordering
cmp forall b. Ord b => b -> Set b -> t b
fs a
q Set a
s = (a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s (t a)) -> t a
forall a b.
(a -> a -> Ordering)
-> (forall s. Reifies s (OrdDict a) => O s b) -> b
withOrd a -> a -> Ordering
cmp ((forall s. Reifies s (OrdDict a) => O s (t a)) -> t a)
-> (forall s. Reifies s (OrdDict a) => O s (t a)) -> t a
forall a b. (a -> b) -> a -> b
$ (Set (O s a) -> t (O s a)) -> Set a -> O s (t a)
forall (f :: * -> *) s a (g :: * -> *).
(f (O s a) -> g (O s a)) -> f a -> O s (g a)
liftOrd1 (O s a -> Set (O s a) -> t (O s a)
forall b. Ord b => b -> Set b -> t b
fs (O s a -> Set (O s a) -> t (O s a))
-> O s a -> Set (O s a) -> t (O s a)
forall a b. (a -> b) -> a -> b
$ a -> O s a
forall s a. a -> O s a
O a
q) Set a
s






-- queryBy'           :: Ord r
--                    => (a -> r)
--                    -> r
--                   -> (forall b. Ord b => b -> Set b -> t b)
--                   -> a -> Set a -> t a
-- queryBy' g fs q s = queryBy
--   where



--   withOrd cmp $ liftOrd1 (fs $ O q) s



  -- withOrd cmp $ liftOrd1 (Set.lookupGE $ O q) s




-- test = queryBy cmpS Set.lookupGE (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]
-- test = succBy cmpS (S "22") $ fromListBy cmpS [S "a" , S "bbb" , S "ddddddd"]