{-# LANGUAGE
  GADTs
  #-}

module Data.Trie.Pred.Disjoint.Tail
  ( DPTrie (..)
  , lookup
  , lookupWithL
  , lookupNearestParent
  , merge
  , areDisjoint
  , litSingletonTail
  , litExtrudeTail
  , sort
  ) where

import Prelude hiding (lookup)
import Data.List.NonEmpty as NE hiding (map, sort)
import Control.Applicative



data DPTrie p t x where
  DMore :: t
        -> Maybe x
        -> [DPTrie p t x]
        -> DPTrie p t x
  DPred :: p
        -> (t -> Maybe r)
        -> Maybe (r -> x)
        -> [DPTrie p t (r -> x)]
        -> DPTrie p t x


-- | Overwrites when similar, leaves untouched when not
merge :: (Eq p, Eq t) => DPTrie p t x -> DPTrie p t x -> DPTrie p t x
merge xx@(DMore t mx xs) yy@(DMore p my ys)
  | t == p = DMore p my $ foldr go [] $ xs ++ ys
  | otherwise = xx
  where
    go :: (Eq p, Eq t) => DPTrie p t x -> [DPTrie p t x] -> [DPTrie p t x]
    go a [] = [a]
    go a (b:bs) | areDisjoint a b =     a : b : bs
                | otherwise       = merge a b : bs
merge xx@(DPred t q mrx xrs) yy@(DPred p w mry yrs)
  | t == p = yy
  | otherwise = xx
merge xx@(DMore t mx xs) yy@(DPred p w mrx xrs) = yy
merge xx@(DPred t q mrx xrs) yy@(DMore p my ys) = yy


areDisjoint :: (Eq p, Eq t) => DPTrie p t x -> DPTrie p t x -> Bool
areDisjoint (DMore t _ _)    (DMore p _ _)    = t == p
areDisjoint (DPred t _ _ _)  (DPred p _ _ _)  = t == p
areDisjoint _ _ = True


lookup :: Eq t => NonEmpty t -> DPTrie p t x -> Maybe x
lookup (t:|ts) (DMore t' mx xs)
  | t == t' = case ts of
    [] -> mx
    _  -> firstJust $ map (lookup $ NE.fromList ts) xs
  | otherwise = Nothing
lookup (t:|ts) (DPred _ p mrx xrs) =
  p t >>=
    \r -> case ts of
      [] -> ($ r) <$> mrx
      _  -> ($ r) <$> firstJust (map (lookup $ NE.fromList ts) xrs)

lookupWithL :: Eq t => (t -> t) -> NonEmpty t -> DPTrie p t x -> Maybe x
lookupWithL f (t:|ts) (DMore t' mx xs)
  | null ts = if f t == t'
              then mx
              else Nothing
  | otherwise = if t == t'
                then firstJust $ map (lookupWithL f $ NE.fromList ts) xs
                else Nothing
lookupWithL f (t:|ts) (DPred _ p mrx xrs) =
  p t >>=
    \r -> case ts of
      [] -> ($ r) <$> mrx
      _  -> ($ r) <$> firstJust (map (lookupWithL f $ NE.fromList ts) xrs)

lookupNearestParent :: Eq t => NonEmpty t -> DPTrie p t x -> Maybe x
lookupNearestParent tss@(t:|ts) trie@(DMore t' mx xs) = case lookup tss trie of
  Nothing -> if t == t'
               then case ts of
                      [] -> mx -- redundant; should have successful lookup
                      _  -> case firstJust $ map (lookupNearestParent $ NE.fromList ts) xs of
                              Nothing -> mx
                              justr   -> justr
               else Nothing
  justr -> justr
lookupNearestParent tss@(t:|ts) trie@(DPred t' p mrx xrs) = case lookup tss trie of
  Nothing -> p t >>=
               \r -> case ts of
                        [] -> ($ r) <$> mrx -- redundant; should have successful lookup
                        _  -> case firstJust $ map (lookupNearestParent $ NE.fromList ts) xrs of
                                Nothing -> ($ r) <$> mrx
                                justr   -> ($ r) <$> justr
  justr -> justr


firstJust :: [Maybe a] -> Maybe a
firstJust [] = Nothing
firstJust (Nothing:xs) = firstJust xs
firstJust (Just x :xs) = Just x


litSingletonTail :: NonEmpty t -> x -> DPTrie p t x
litSingletonTail (t:|[]) x = DMore t (Just x) []
litSingletonTail (t:|ts) x = DMore t Nothing  [litSingletonTail (NE.fromList ts) x]


litExtrudeTail :: [t] -> DPTrie p t x -> DPTrie p t x
litExtrudeTail [] r = r
litExtrudeTail (t:ts) r = DMore t Nothing [litExtrudeTail ts r]


sort :: (Eq p, Eq t) => [DPTrie p t x] -> [DPTrie p t x]
sort = foldr insert []
  where
    insert :: (Eq p, Eq t) => DPTrie p t x -> [DPTrie p t x] -> [DPTrie p t x]
    insert r [] = [r]
    insert x@(DMore t _ _) (y@(DMore p _ _):rs)
      | t == p = x : rs
      | otherwise = x : y : rs
    insert x@(DMore t _ _) (y@(DPred p _ _ _):rs) =
        x : y : rs
    insert x@(DPred t _ _ _) (y@(DPred p _ _ _):rs)
      | t == p = x : rs -- basis
      | otherwise = x : y : rs
    insert x@(DPred t _ _ _) (y@(DMore p _ _):rs) =
        y : insert x rs