{-# LANGUAGE
      CPP,
      DerivingVia
  #-}

module Data.Mapping.Piecewise where

#if MIN_VERSION_GLASGOW_HASKELL(9,6,0,0)
#else
import Control.Applicative (liftA2)
#endif
import Control.Applicative (liftA3)
import Data.Algebra.Boolean
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Data.Mapping


-- | A data structure storing mappings that are constant on
-- intervals.
--
-- If the space of keys not discrete, then these mappings are
-- right-continuous: values are in general defined on intervals $a
-- \leq x < b$ which are closed on the left and open on the right.
data Piecewise k v = Piecewise {
  forall k v. Piecewise k v -> v
leftEnd :: v,
  forall k v. Piecewise k v -> Map k v
starts :: Map k v
} deriving (Piecewise k v -> Piecewise k v -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
/= :: Piecewise k v -> Piecewise k v -> Bool
$c/= :: forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
== :: Piecewise k v -> Piecewise k v -> Bool
$c== :: forall k v. (Eq v, Eq k) => Piecewise k v -> Piecewise k v -> Bool
Eq, Piecewise k v -> Piecewise k v -> Bool
Piecewise k v -> Piecewise k v -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {k} {v}. (Ord v, Ord k) => Eq (Piecewise k v)
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Ordering
forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
min :: Piecewise k v -> Piecewise k v -> Piecewise k v
$cmin :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
max :: Piecewise k v -> Piecewise k v -> Piecewise k v
$cmax :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Piecewise k v
>= :: Piecewise k v -> Piecewise k v -> Bool
$c>= :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
> :: Piecewise k v -> Piecewise k v -> Bool
$c> :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
<= :: Piecewise k v -> Piecewise k v -> Bool
$c<= :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
< :: Piecewise k v -> Piecewise k v -> Bool
$c< :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Bool
compare :: Piecewise k v -> Piecewise k v -> Ordering
$ccompare :: forall k v.
(Ord v, Ord k) =>
Piecewise k v -> Piecewise k v -> Ordering
Ord)

piecewiseFromAsc :: Eq k => v -> [(k,v)] -> Piecewise k v
piecewiseFromAsc :: forall k v. Eq k => v -> [(k, v)] -> Piecewise k v
piecewiseFromAsc v
k = forall k v. v -> Map k v -> Piecewise k v
Piecewise v
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Eq k => [(k, a)] -> Map k a
M.fromAscList

instance (Show k, Show v) => Show (Piecewise k v) where
  showsPrec :: Int -> Piecewise k v -> ShowS
showsPrec Int
d (Piecewise v
k Map k v
m) =
    (String
"piecewiseFromAsc " <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a. Show a => Int -> a -> ShowS
showsPrec Int
d v
k forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    (String
" " <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a. Show a => [a] -> ShowS
showList (forall k a. Map k a -> [(k, a)]
M.toList Map k v
m)

changeAt :: v -> k -> v -> Piecewise k v
changeAt :: forall v k. v -> k -> v -> Piecewise k v
changeAt v
a k
x v
b = forall k v. v -> Map k v -> Piecewise k v
Piecewise v
a forall a b. (a -> b) -> a -> b
$ forall k a. k -> a -> Map k a
M.singleton k
x v
b

atLeast :: k -> Piecewise k Bool
atLeast :: forall k. k -> Piecewise k Bool
atLeast k
k = forall v k. v -> k -> v -> Piecewise k v
changeAt Bool
False k
k Bool
True

lessThan :: k -> Piecewise k Bool
lessThan :: forall k. k -> Piecewise k Bool
lessThan k
k = forall v k. v -> k -> v -> Piecewise k v
changeAt Bool
True k
k Bool
False

fromAscList :: (Ord k, Eq v) => v -> [(k,v)] -> Piecewise k v
fromAscList :: forall k v. (Ord k, Eq v) => v -> [(k, v)] -> Piecewise k v
fromAscList = let
  inner :: b -> [(a, b)] -> [(a, b)]
inner b
_ [] = []
  inner b
a ((a
y,b
b):[(a, b)]
r)
    | b
a forall a. Eq a => a -> a -> Bool
== b
b    = b -> [(a, b)] -> [(a, b)]
inner b
a [(a, b)]
r
    | Bool
otherwise = (a
y,b
b)forall a. a -> [a] -> [a]
:b -> [(a, b)] -> [(a, b)]
inner b
b [(a, b)]
r
  run :: v -> [(k, v)] -> Piecewise k v
run v
x = forall k v. v -> Map k v -> Piecewise k v
Piecewise v
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Eq k => [(k, a)] -> Map k a
M.fromAscList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {b} {a}. Eq b => b -> [(a, b)] -> [(a, b)]
inner v
x
  in forall {k} {v}. (Eq k, Eq v) => v -> [(k, v)] -> Piecewise k v
run

values :: Piecewise k v -> [v]
values :: forall k v. Piecewise k v -> [v]
values (Piecewise v
x Map k v
m) = v
x forall a. a -> [a] -> [a]
: forall k a. Map k a -> [a]
M.elems Map k v
m

instance Foldable (Piecewise k) where
  foldMap :: forall m a. Monoid m => (a -> m) -> Piecewise k a -> m
foldMap a -> m
f (Piecewise a
a Map k a
m) = a -> m
f a
a forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Map k a
m

instance Ord k => Mapping k (Piecewise k) where

  cst :: forall v. v -> Piecewise k v
cst v
x = forall k v. v -> Map k v -> Piecewise k v
Piecewise v
x forall k a. Map k a
M.empty

  act :: forall v. Piecewise k v -> k -> v
act (Piecewise v
a Map k v
f) k
x = case forall k v. Ord k => k -> Map k v -> Maybe (k, v)
M.lookupLE k
x Map k v
f of
    Maybe (k, v)
Nothing -> v
a
    Just (k
_,v
b) -> v
b

  isConst :: forall v. Ord v => Piecewise k v -> Maybe v
isConst (Piecewise v
a Map k v
f) = if forall k a. Map k a -> Bool
M.null Map k v
f then forall a. a -> Maybe a
Just v
a else forall a. Maybe a
Nothing

  mmap :: forall v u. Ord v => (u -> v) -> Piecewise k u -> Piecewise k v
mmap u -> v
p (Piecewise u
a Map k u
f) = forall k v. (Ord k, Eq v) => v -> [(k, v)] -> Piecewise k v
fromAscList (u -> v
p u
a) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap u -> v
p forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k a. Map k a -> [(k, a)]
M.toList Map k u
f)

  mtraverse :: forall (f :: * -> *) v u.
(Applicative f, Ord v) =>
(u -> f v) -> Piecewise k u -> f (Piecewise k v)
mtraverse u -> f v
p (Piecewise u
a Map k u
f) = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall k v. (Ord k, Eq v) => v -> [(k, v)] -> Piecewise k v
fromAscList (u -> f v
p u
a) (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse u -> f v
p) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Map k u
f)

  merge :: forall w u v.
Ord w =>
(u -> v -> w) -> Piecewise k u -> Piecewise k v -> Piecewise k w
merge u -> v -> w
p = let

    inner :: u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b w
c r :: [(a, u)]
r@((a
x,u
a'):[(a, u)]
r') s :: [(a, v)]
s@((a
y,v
b'):[(a, v)]
s') = case forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
      Ordering
LT -> let
        c' :: w
c' = u -> v -> w
p u
a' v
b
        in if w
c' forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c [(a, u)]
r' [(a, v)]
s else (a
x,w
c')forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c' [(a, u)]
r' [(a, v)]
s
      Ordering
GT -> let
        c' :: w
c' = u -> v -> w
p u
a v
b'
        in if w
c' forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c [(a, u)]
r [(a, v)]
s' else (a
y,w
c')forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c' [(a, u)]
r [(a, v)]
s'
      Ordering
EQ -> let
        c' :: w
c' = u -> v -> w
p u
a' v
b'
        in if w
c' forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b' w
c [(a, u)]
r' [(a, v)]
s' else (a
x,w
c')forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b' w
c' [(a, u)]
r' [(a, v)]
s'
    inner u
a v
_ w
c [] ((a
y,v
b'):[(a, v)]
s') = let
      c' :: w
c' = u -> v -> w
p u
a v
b'
      in if w
c' forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c [] [(a, v)]
s' else (a
y,w
c')forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b' w
c' [] [(a, v)]
s'
    inner u
_ v
b w
c ((a
x,u
a'):[(a, u)]
r') [] = let
      c' :: w
c' = u -> v -> w
p u
a' v
b
      in if w
c' forall a. Eq a => a -> a -> Bool
== w
c then u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c [(a, u)]
r' [] else (a
x,w
c')forall a. a -> [a] -> [a]
:u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a' v
b w
c' [(a, u)]
r' []
    inner u
_ v
_ w
_ [] [] = []

    run :: Piecewise k u -> Piecewise k v -> Piecewise k w
run (Piecewise u
a Map k u
f) (Piecewise v
b Map k v
g) = let
      c :: w
c = u -> v -> w
p u
a v
b
      l :: [(k, w)]
l = forall {a}.
Ord a =>
u -> v -> w -> [(a, u)] -> [(a, v)] -> [(a, w)]
inner u
a v
b w
c (forall k a. Map k a -> [(k, a)]
M.toList Map k u
f) (forall k a. Map k a -> [(k, a)]
M.toList Map k v
g)
      in forall k v. v -> Map k v -> Piecewise k v
Piecewise w
c forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k, w)]
l

    in forall {k}.
Ord k =>
Piecewise k u -> Piecewise k v -> Piecewise k w
run

  mergeA :: forall (f :: * -> *) w u v.
(Applicative f, Ord w) =>
(u -> v -> f w)
-> Piecewise k u -> Piecewise k v -> f (Piecewise k w)
mergeA u -> v -> f w
p = let

    maybePrepend :: a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x b
u b
v [(a, b)]
l
      | b
u forall a. Eq a => a -> a -> Bool
== b
v    = [(a, b)]
l
      | Bool
otherwise = (a
x,b
v)forall a. a -> [a] -> [a]
:[(a, b)]
l

    inner :: u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b f w
c r :: [(a, u)]
r@((a
x,u
a'):[(a, u)]
r') s :: [(a, v)]
s@((a
y,v
b'):[(a, v)]
s') = case forall a. Ord a => a -> a -> Ordering
compare a
x a
y of
      Ordering
LT -> let
        c' :: f w
c' = u -> v -> f w
p u
a' v
b
        in forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b f w
c' [(a, u)]
r' [(a, v)]
s
      Ordering
GT -> let
        c' :: f w
c' = u -> v -> f w
p u
a v
b'
        in forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
y) f w
c f w
c' forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b' f w
c' [(a, u)]
r [(a, v)]
s'
      Ordering
EQ -> let
        c' :: f w
c' = u -> v -> f w
p u
a' v
b'
        in forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b' f w
c' [(a, u)]
r' [(a, v)]
s'
    inner u
a v
_ f w
c [] ((a
y,v
b'):[(a, v)]
s') = let
      c' :: f w
c' = u -> v -> f w
p u
a v
b'
      in forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
y) f w
c f w
c' forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b' f w
c' [] [(a, v)]
s'
    inner u
_ v
b f w
c ((a
x,u
a'):[(a, u)]
r') [] = let
      c' :: f w
c' = u -> v -> f w
p u
a' v
b
      in forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (forall {b} {a}. Eq b => a -> b -> b -> [(a, b)] -> [(a, b)]
maybePrepend a
x) f w
c f w
c' forall a b. (a -> b) -> a -> b
$ u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a' v
b f w
c' [(a, u)]
r' []
    inner u
_ v
_ f w
_ [] [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []

    run :: Piecewise k u -> Piecewise k v -> f (Piecewise k w)
run (Piecewise u
a Map k u
f) (Piecewise v
b Map k v
g) = let
      c :: f w
c = u -> v -> f w
p u
a v
b
      l :: f [(k, w)]
l = forall {a}.
Ord a =>
u -> v -> f w -> [(a, u)] -> [(a, v)] -> f [(a, w)]
inner u
a v
b f w
c (forall k a. Map k a -> [(k, a)]
M.toList Map k u
f) (forall k a. Map k a -> [(k, a)]
M.toList Map k v
g)
      in forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall k v. v -> Map k v -> Piecewise k v
Piecewise f w
c (forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f [(k, w)]
l)

    in forall {k}.
Ord k =>
Piecewise k u -> Piecewise k v -> f (Piecewise k w)
run

instance Neighbourly (Piecewise k) where
  neighbours :: forall v. Ord v => Piecewise k v -> Set (v, v)
neighbours Piecewise k v
m = let
    v :: [v]
v = forall k v. Piecewise k v -> [v]
values Piecewise k v
m
    in forall a. Ord a => [a] -> Set a
S.fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [v]
v (forall a. [a] -> [a]
tail [v]
v)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Semigroup b) => Semigroup (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Monoid b) => Monoid (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Num b) => Num (Piecewise k b)

deriving via (AlgebraWrapper k (Piecewise k) b)
  instance (Ord k, Ord b, Boolean b) => Boolean (Piecewise k b)