--------------------------------------------------------------------------------
-- |
-- Module      :  HGeometry.Trie.Type
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--
-- A Trie type
--
--------------------------------------------------------------------------------
module HGeometry.Trie.Type
  ( TrieF(..)
  , root
  , mapWithEdgeLabels
  , foldWithEdgeLabels
  ) where

import           Control.Lens
import           Data.Bifoldable
import           Data.Bitraversable
import           Data.Foldable1
import           Data.Functor.Apply ((<.*>))
import           Data.Functor.Classes
import           Data.Semigroup.Traversable

-- import Debug.Trace
--------------------------------------------------------------------------------

-- | The Trie data type, parameterized by the data structure storing the children.
data TrieF f e v = Node v (f e (TrieF f e v))

--  | Access the root of the trie
root :: Lens' (TrieF f e v) v
root :: forall {k} (f :: k -> * -> *) (e :: k) v (f :: * -> *).
Functor f =>
(v -> f v) -> TrieF f e v -> f (TrieF f e v)
root = (TrieF f e v -> v)
-> (TrieF f e v -> v -> TrieF f e v)
-> Lens (TrieF f e v) (TrieF f e v) v v
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\(Node v
x f e (TrieF f e v)
_) -> v
x) (\(Node v
_ f e (TrieF f e v)
chs) v
x -> v -> f e (TrieF f e v) -> TrieF f e v
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node v
x f e (TrieF f e v)
chs)
{-# INLINE root #-}


deriving instance (Show v, Show e, Show2 f) => Show (TrieF f e v)
deriving instance (Eq v, Eq e, Eq2 f)       => Eq (TrieF f e v)
deriving instance (Ord v, Ord e, Ord2 f)    => Ord (TrieF f e v)

deriving instance (Functor (f e))     => Functor (TrieF f e)
deriving instance (Foldable (f e))    => Foldable (TrieF f e)
deriving instance (Traversable (f e)) => Traversable (TrieF f e)

-- instance Foldable (f e) => Foldable1 (TrieF f e) where
--   foldMap1 f (Node v chs) = let Endo g = foldMap (\x -> Endo $ \x0 -> x0 <> foldMap1 f x) chs
--                             in g (f v)
-- somehow the order is wrong here...

instance Foldable (f e) => Foldable1 (TrieF f e) where
  foldMap1 :: forall m a. Semigroup m => (a -> m) -> TrieF f e a -> m
foldMap1 a -> m
f = TrieF f e a -> m
go
    where
      go :: TrieF f e a -> m
go (Node a
v f e (TrieF f e a)
chs) = case (TrieF f e a -> Maybe m) -> f e (TrieF f e a) -> Maybe m
forall m a. Monoid m => (a -> m) -> f e a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (m -> Maybe m
forall a. a -> Maybe a
Just (m -> Maybe m) -> (TrieF f e a -> m) -> TrieF f e a -> Maybe m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TrieF f e a -> m
go) f e (TrieF f e a)
chs of
                          Maybe m
Nothing -> a -> m
f a
v
                          Just m
s  -> a -> m
f a
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
s

  -- f (Node v chs) = let Endo g = foldMap (\x -> Endo $ \x0 -> x0 <> foldMap1 f x) chs
  --                           in g (f v)


instance Traversable (f e) => Traversable1 (TrieF f e) where
  traverse1 :: forall (f :: * -> *) a b.
Apply f =>
(a -> f b) -> TrieF f e a -> f (TrieF f e b)
traverse1 a -> f b
f = TrieF f e a -> f (TrieF f e b)
go
    where
      go :: TrieF f e a -> f (TrieF f e b)
go (Node a
v f e (TrieF f e a)
chs) = b -> f e (TrieF f e b) -> TrieF f e b
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node (b -> f e (TrieF f e b) -> TrieF f e b)
-> f b -> f (f e (TrieF f e b) -> TrieF f e b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
v f (f e (TrieF f e b) -> TrieF f e b)
-> MaybeApply f (f e (TrieF f e b)) -> f (TrieF f e b)
forall (f :: * -> *) a b.
Apply f =>
f (a -> b) -> MaybeApply f a -> f b
<.*> (TrieF f e a -> f (TrieF f e b))
-> f e (TrieF f e a) -> MaybeApply f (f e (TrieF f e b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Apply f) =>
(a -> f b) -> t a -> MaybeApply f (t b)
traverse1Maybe TrieF f e a -> f (TrieF f e b)
go f e (TrieF f e a)
chs

instance Bifunctor f => Bifunctor (TrieF f) where
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> TrieF f a c -> TrieF f b d
bimap a -> b
f c -> d
g = TrieF f a c -> TrieF f b d
go
    where
      go :: TrieF f a c -> TrieF f b d
go (Node c
v f a (TrieF f a c)
chs) = d -> f b (TrieF f b d) -> TrieF f b d
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node (c -> d
g c
v) ((a -> b)
-> (TrieF f a c -> TrieF f b d)
-> f a (TrieF f a c)
-> f b (TrieF f b d)
forall a b c d. (a -> b) -> (c -> d) -> f a c -> f b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f TrieF f a c -> TrieF f b d
go f a (TrieF f a c)
chs)

instance Bifoldable f => Bifoldable (TrieF f) where
  bifoldMap :: forall m a b. Monoid m => (a -> m) -> (b -> m) -> TrieF f a b -> m
bifoldMap a -> m
f b -> m
g = TrieF f a b -> m
go
    where
      go :: TrieF f a b -> m
go (Node b
v f a (TrieF f a b)
chs) = b -> m
g b
v m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (a -> m) -> (TrieF f a b -> m) -> f a (TrieF f a b) -> m
forall m a b. Monoid m => (a -> m) -> (b -> m) -> f a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap a -> m
f TrieF f a b -> m
go f a (TrieF f a b)
chs

instance Bitraversable f => Bitraversable (TrieF f) where
  bitraverse :: forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> TrieF f a b -> f (TrieF f c d)
bitraverse a -> f c
f b -> f d
g = TrieF f a b -> f (TrieF f c d)
go
    where
      go :: TrieF f a b -> f (TrieF f c d)
go (Node b
v f a (TrieF f a b)
chs) = d -> f c (TrieF f c d) -> TrieF f c d
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node (d -> f c (TrieF f c d) -> TrieF f c d)
-> f d -> f (f c (TrieF f c d) -> TrieF f c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> b -> f d
g b
v f (f c (TrieF f c d) -> TrieF f c d)
-> f (f c (TrieF f c d)) -> f (TrieF f c d)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f c)
-> (TrieF f a b -> f (TrieF f c d))
-> f a (TrieF f a b)
-> f (f c (TrieF f c d))
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> f a b -> f (f c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse a -> f c
f TrieF f a b -> f (TrieF f c d)
go f a (TrieF f a b)
chs



-- | A mapping function to transform the values of the trie, possibly using the edge
-- labels.  note that the root does not have an edge label, so it is transformed using a
-- separate function.
mapWithEdgeLabels                          :: FunctorWithIndex e (f e)
                                           => (v -> v')
                                           -- ^ function by which to transform the root
                                           -> (e -> v -> v')
                                           -> TrieF f e v -> TrieF f e v'
mapWithEdgeLabels :: forall e (f :: * -> * -> *) v v'.
FunctorWithIndex e (f e) =>
(v -> v') -> (e -> v -> v') -> TrieF f e v -> TrieF f e v'
mapWithEdgeLabels v -> v'
fRoot e -> v -> v'
f (Node v
root' f e (TrieF f e v)
chs) = v' -> f e (TrieF f e v') -> TrieF f e v'
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node (v -> v'
fRoot v
root') ((e -> TrieF f e v -> TrieF f e v')
-> f e (TrieF f e v) -> f e (TrieF f e v')
forall a b. (e -> a -> b) -> f e a -> f e b
forall i (f :: * -> *) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap e -> TrieF f e v -> TrieF f e v'
go f e (TrieF f e v)
chs)
  where
    go :: e -> TrieF f e v -> TrieF f e v'
go e
e (Node v
x f e (TrieF f e v)
chs') = v' -> f e (TrieF f e v') -> TrieF f e v'
forall {k} (f :: k -> * -> *) (e :: k) v.
v -> f e (TrieF f e v) -> TrieF f e v
Node (e -> v -> v'
f e
e v
x) ((e -> TrieF f e v -> TrieF f e v')
-> f e (TrieF f e v) -> f e (TrieF f e v')
forall a b. (e -> a -> b) -> f e a -> f e b
forall i (f :: * -> *) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap e -> TrieF f e v -> TrieF f e v'
go f e (TrieF f e v)
chs')

-- | fold with the edge labels
foldWithEdgeLabels :: (FoldableWithIndex e (f e), Monoid m)
                   => (v -> m) -- ^ function by which to transform the root
                   -> (e -> v -> m)
                   -> TrieF f e v
                   -> m
foldWithEdgeLabels :: forall e (f :: * -> * -> *) m v.
(FoldableWithIndex e (f e), Monoid m) =>
(v -> m) -> (e -> v -> m) -> TrieF f e v -> m
foldWithEdgeLabels v -> m
fRoot e -> v -> m
f (Node v
root' f e (TrieF f e v)
chs) = v -> m
fRoot v
root' m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (e -> TrieF f e v -> m) -> f e (TrieF f e v) -> m
forall m a. Monoid m => (e -> a -> m) -> f e a -> m
forall i (f :: * -> *) m a.
(FoldableWithIndex i f, Monoid m) =>
(i -> a -> m) -> f a -> m
ifoldMap e -> TrieF f e v -> m
go f e (TrieF f e v)
chs
  where
    go :: e -> TrieF f e v -> m
go e
e (Node v
x f e (TrieF f e v)
chs') = e -> v -> m
f e
e v
x m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (e -> TrieF f e v -> m) -> f e (TrieF f e v) -> m
forall m a. Monoid m => (e -> a -> m) -> f e a -> m
forall i (f :: * -> *) m a.
(FoldableWithIndex i f, Monoid m) =>
(i -> a -> m) -> f a -> m
ifoldMap e -> TrieF f e v -> m
go f e (TrieF f e v)
chs'

-- instance Functor f => FunctorWithIndex e (TrieF f e) where
--   imap f = go
--     where
--       go (Node x chs) = Node


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

{-

-- | A binary Trie type
type BinaryTrie e v = TrieF (KV AtMostTwo) e v

-- | Pattern match on a leaf node
pattern Leaf   :: v -> BinaryTrie e v
pattern Leaf v = Node v (KV Zero)

-- | Pattern match on a node with one child
pattern OneNode     :: v -> (e, BinaryTrie e v) -> BinaryTrie e v
pattern OneNode v x = Node v (KV (One x))

-- | Pattern match on a node with two children
pattern TwoNode      :: v -> (e, BinaryTrie e v) -> (e, BinaryTrie e v) -> BinaryTrie e v
pattern TwoNode v l r = Node v (KV (Two l r))

{-# COMPLETE Leaf, OneNode, TwoNode #-}

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

-- | Trie to convert the trie into a binary trie.
asBinaryTrie              :: Traversable f => TrieF (KV f) e v -> Maybe (BinaryTrie e v)
asBinaryTrie (Node x chs) = traverse asBinaryTrie chs >>= \res -> case F.toList (assocs res) of
                              []    -> pure $ Leaf x
                              [c]   -> pure $ OneNode x c
                              [l,r] -> pure $ TwoNode x l r
                              _     -> Nothing


-}