{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module KMC.SymbolicFST
       (FST(..)
       ,Edge
       ,fstStateSize
       ,fstTransSize
       ,edgesToList
       ,edgesFromList
       ,fstEvalEpsilonEdges
       ,fstEdges
       ,fstAbstractEvalEdgesAll
       ,isChoiceState
       ,isSkipState
       ,isJoinState
       ,coarsestPredicateSet
       ,prefixTests
       ,rightClosure,rightInputClosure
       ,mapEdges
       ,trim
       ,eForward,eBackward,eForwardEpsilon,eBackwardEpsilon
       ,enumerateStates
       ,run
       ,runSequential
       ,runBacktracking
       ) where

import           Control.Applicative
import           Control.Monad (guard)
import           Data.Map ((!))
import qualified Data.Map as M
import           Data.Maybe (maybeToList)
import           Data.Monoid
import qualified Data.Set as S

import           KMC.Backtracking
import           KMC.Theories

import           Prelude

-- | (Non-)deterministic Finite State Transducer. To be well-formed, the
-- following conditions must be met:
--
-- (i) For any state s, the only transitions out of s are epsilon-transitions or
-- symbol transitions, but not both.
data FST q pred func =
  FST
  { fstS :: S.Set q
  , fstE :: OrderedEdgeSet q pred func
  , fstI :: q
  , fstF :: S.Set q
  }

-- | Type synonym for a transition. A transition is either labeled by a
-- predicate and an output function; or it is labeled by an output string. This
-- models a set of transitions between the two states where the output label may
-- depend on the symbol being read; in the case of epsilon-transitions, a plain
-- output string is provided.
type Edge q pred func = (q, Either (pred, func) (Rng func), q)

-- | An ordered edge set with fast reverse lookup.
data OrderedEdgeSet q pred func =
  OrderedEdgeSet { eForward :: M.Map q [(pred,func,q)]
                 , eBackward :: M.Map q [(pred,func,q)]
                 , eForwardEpsilon :: M.Map q [(Rng func, q)]
                 , eBackwardEpsilon :: M.Map q [(Rng func, q)]
                 }

deriving instance (Eq q, Eq pred, Eq func, Eq (Rng func)) => Eq (OrderedEdgeSet q pred func)
deriving instance (Ord q, Ord pred, Ord func, Ord (Rng func)) => Ord (OrderedEdgeSet q pred func)
deriving instance (Show q, Show pred, Show func, Show (Rng func)) => Show (OrderedEdgeSet q pred func)
deriving instance (Eq q, Eq pred, Eq func, Eq (Rng func)) => Eq (FST q pred func)
deriving instance (Ord q, Ord pred, Ord func, Ord (Rng func)) => Ord (FST q pred func)
deriving instance (Show q, Show pred, Show func, Show (Rng func)) => Show (FST q pred func)

-- | Construct an edge set from a list of edges
edgesFromList :: (Ord q) => [Edge q pred func] -> OrderedEdgeSet q pred func
edgesFromList es =
  OrderedEdgeSet
  { eForward = M.fromListWith (flip (++)) [ (q, [(a, b, q')]) | (q, Left (a,b), q') <- es ]
  , eBackward = M.fromListWith (flip (++)) [ (q', [(a, b, q)]) | (q, Left (a,b), q') <- es ]
  , eForwardEpsilon = M.fromListWith (flip (++)) [(q, [(y, q')]) | (q, Right y, q') <- es ]
  , eBackwardEpsilon = M.fromListWith (flip (++)) [(q', [(y,q)]) | (q, Right y, q') <- es ]
  }

-- | Get the number of states in the FST
fstStateSize :: FST q pred func -> Int
fstStateSize = S.size . fstS

-- | Get the number of transitions in the FST
fstTransSize :: FST q pred func -> Int
fstTransSize fst' = M.size (eForward (fstE fst')) + M.size (eForwardEpsilon (fstE fst'))

{-
unionEdges :: (Ord q) => OrderedEdgeSet q pred func
           -> OrderedEdgeSet q pred func
           -> OrderedEdgeSet q pred func
unionEdges es es' =
  OrderedEdgeSet
  { eForward  = M.unionWith (++) (eForward es) (eForward es')
  , eBackward = M.unionWith (++) (eBackward es) (eBackward es')
  , eForwardEpsilon = M.unionWith (++) (eForwardEpsilon es) (eForwardEpsilon es')
  , eBackwardEpsilon = M.unionWith (++) (eBackwardEpsilon es) (eBackwardEpsilon es')
  }
-}

flipEdges :: OrderedEdgeSet q pred func -> OrderedEdgeSet q pred func
flipEdges es = OrderedEdgeSet {eForward = eBackward es
                              ,eBackward = eForward es
                              ,eForwardEpsilon = eBackwardEpsilon es
                              ,eBackwardEpsilon = eForwardEpsilon es}

-- | Get the underlying list representation of an ordered edge set
edgesToList :: OrderedEdgeSet q pred func -> [(q, Either (pred, func) (Rng func), q)]
edgesToList es = symEdgesToList (eForward es) ++ epsEdgesToList (eForwardEpsilon es)

-- | Get the list representation of only the symbol edges in an ordered edge set
symEdgesToList :: M.Map q [(pred, func, q)] -> [(q, Either (pred, func) (Rng func), q)]
symEdgesToList es = [ (q, Left (a,b), q') | (q, xs) <- M.toList es, (a,b,q') <- xs ]

-- | Get the list representation of only the epsilon edges in an ordered edge set
epsEdgesToList :: M.Map q [(Rng func, q)] -> [(q, Either (pred, func) (Rng func), q)]
epsEdgesToList es = [ (q, Right y, q') | (q, xs) <- M.toList es, (y, q') <- xs ]

-- | Map four functions over the edges in an FST. Symbol edges may be turned
-- into epsilon edges and vice versa.
mapEdges :: (Ord q)
         => (q -> [(pred1,func1,q)] -> [(pred2,func2,q)])
         -> (q -> [(pred1,func1,q)] -> [(Rng func2, q)])
         -> (q -> [(Rng func1, q)] -> [(pred2,func2,q)])
         -> (q -> [(Rng func1, q)] -> [(Rng func2, q)])
         -> FST q pred1 func1 -> FST q pred2 func2
mapEdges symsym symeps epssym epseps fst' =
  fst' { fstE = edgesFromList (symEdges1 ++ symEdges2 ++ epsEdges1 ++ epsEdges2) }
  where
    esym = eForward $ fstE fst'
    eeps = eForwardEpsilon $ fstE fst'
    symEdges1 = symEdgesToList $ M.mapWithKey symsym $ esym
    symEdges2 = symEdgesToList $ M.mapWithKey epssym $ eeps
    epsEdges1 = epsEdgesToList $ M.mapWithKey symeps $ esym
    epsEdges2 = epsEdgesToList $ M.mapWithKey epseps $ eeps

-- | Given an input symbol and a starting state, compute the concrete set of
-- transitions from that state.
evalEdges :: (Ord st
             ,Function func
             ,SetLike pred (Dom func))
          => OrderedEdgeSet st pred func
          -> st -> Dom func -> [(Rng func, st)]
evalEdges (OrderedEdgeSet { eForward = me }) q x =
  case M.lookup q me of
    Nothing -> []
    Just es -> concatMap evalEdge es
      where
        evalEdge (p, f, q')
          | member x p = [(eval f x, q')]
          | otherwise = []

-- | Given a state, compute the concrete set of epsilon-transitions from that
-- state.
evalEpsilonEdges :: (Ord st) => OrderedEdgeSet st pred func
                 -> st -> [(Rng func, st)]
evalEpsilonEdges (OrderedEdgeSet { eForwardEpsilon = meps }) q =
  case M.lookup q meps of
    Nothing -> []
    Just es -> es

-- | If abstractEvalEdgesAll fst q phi == [(f1, q1), ..., (fn, qn)], then
--   for all qi and for all a in [[phi]], q steps to qi reading a.
abstractEvalEdgesAll :: (Ord st
                        ,PartialOrder pred)
                       =>
                       OrderedEdgeSet st pred func
                     -> st -> pred -> [(func, st)]
abstractEvalEdgesAll (OrderedEdgeSet { eForward = me}) q p =
  case M.lookup q me of
    Nothing -> []
    Just es -> [ (f, q') | (p', f, q') <- es, p `lte` p' ]

-- | Like evalEpsilonEdges, but given an FST.
fstEvalEpsilonEdges :: (Ord st) => FST st pred func -> st -> [(Rng func, st)]
fstEvalEpsilonEdges aut = evalEpsilonEdges (fstE aut)

-- | Get the list of forward symbol edges from a state
fstEdges :: (Ord st) => FST st pred func -> st -> [(pred, func, st)]
fstEdges fst' q =
  case M.lookup q (eForward $ fstE fst') of
    Nothing -> []
    Just es -> es

-- | Like fstEvalEdges, but given an FST.
fstEvalEdges :: (Ord st
                ,Function func
                ,SetLike pred (Dom func))
                => FST st pred func -> st -> Dom func -> [(Rng func, st)]
fstEvalEdges fst' q a = evalEdges (fstE fst') q a

-- | Like abstractEvalEdgesAll, but given an FST.
fstAbstractEvalEdgesAll :: (Ord st
                           ,PartialOrder pred)
                          => FST st pred func -> st -> pred -> [(func, st)]
fstAbstractEvalEdgesAll aut = abstractEvalEdgesAll (fstE aut)

-- | Is the given state a non-deterministic choice state, or an input action
-- state?
isChoiceState :: (Ord st) => FST st pred func -> st -> Bool
isChoiceState fst' q =
  case (M.member q (eForward . fstE $ fst'), M.member q (eForwardEpsilon . fstE $ fst')) of
    (True, True) -> error "Inconsistent FST - a state is both a choice and symbol state"
    (_, b)       -> b

isSkipState :: (Ord st) => FST st pred func -> st -> Bool
isSkipState fst' q =
  case M.lookup q (eForwardEpsilon . fstE $ fst') of
    Just [_] -> True
    _ -> False

isJoinState :: (Ord st) => FST st pred func -> st -> Bool
isJoinState fst' q =
  let nBackEdges = length (concat (maybeToList (M.lookup q (eBackwardEpsilon . fstE $ fst'))))
                 + length (concat (maybeToList (M.lookup q (eBackward . fstE $ fst'))))
  in nBackEdges > 1

-- | Given a state set A (represented as a list), compute the coarsest predicate
-- set obtained from the set
-- { p | p is a predicate on a transition starting in q, q in A }
coarsestPredicateSet :: (Boolean pred
                        ,PartialOrder pred
                        ,Ord st
                        ,Ord pred) =>
                        FST st pred func
                     -> [st]
                     -> [pred]
coarsestPredicateSet fst' qs = coarsestPartition ps
  where
    ps = S.toList $ S.fromList
           [ p | q <- qs
               , (p, _, _) <- maybe [] id (M.lookup q (eForward . fstE $ fst')) ]

-- | Compute an ordered right closure with output
rightClosure :: (Ord st, Monoid (Rng func)) => FST st pred func -> st -> [(Rng func, st)]
rightClosure fst' = snd . go S.empty mempty
    where
      go vis out q =
        case fstEvalEpsilonEdges fst' q of
          [] -> (vis, [(out, q)])
          xs -> foldl (\(vis', acc) (w, q') ->
                        if S.member q' vis' then
                            (vis', acc)
                        else
                            let (vis'', ys) = go (S.insert q' vis') (mappend out w) q'
                            in (vis'', acc ++ ys))
                      (vis, [])
                      xs

---------------
-- LCP analysis
---------------

-- | Compute the longest deterministic prefix for a pointed state set
-- (A, q).
ldp :: (Ord st, Ord pred, PartialOrder pred, Boolean pred) =>
       FST st pred func -> S.Set st -> st -> [pred]
ldp fst' = go
    where
      go ctx q =
          case M.lookup q (eForward $ fstE fst') of
            Just [(p,_,q')] | p `elem` (coarsestPredicateSet fst' $ S.toList ctx) ->
                                p:go (stepAll fst' p ctx) q'
            _ -> []

-- | Given a set A and a predicate p, compute the set A' of states that can be reached via a
-- transition from A labeled by a predicate containing p.
stepAll :: (Ord st, PartialOrder pred) => FST st pred func -> pred -> S.Set st -> S.Set st
stepAll fst' p = S.unions . map aux . S.toList
    where
      aux q = S.unions $ map (rightInputClosure fst' . snd)
                       $ fstAbstractEvalEdgesAll fst' q p

-- | Compute an unordered right closure without output on the input automaton of an FST
rightInputClosure :: (Ord st) => FST st pred func -> st -> S.Set st
rightInputClosure fst' = snd . go S.empty
    where
      go vis q =
       case fstEvalEpsilonEdges fst' q of
         [] -> (vis, S.singleton q)
         xs -> foldl (\(vis', acc) (_, q') ->
                       if S.member q' vis then
                         (vis', acc)
                       else
                         let (vis'', ys) = go (S.insert q' vis') q'
                         in (vis'', S.union acc ys))
                     (vis, S.empty)
                     xs

prefixTests :: (Boolean pred, PartialOrder pred, Ord st, Ord pred) =>
               FST st pred func
            -> Bool
            -> [st]
            -> [([pred], S.Set st)]
prefixTests fst' singletonMode states =
  [ (t, killed t) | t <- tests ]
  where
    ldps = [ (ldp fst' (S.fromList states) q, q) | not singletonMode, q <- states ]
    tests = S.toList $ S.fromList $ [ [p] | p <- coarsestPredicateSet fst' states ] ++ map fst ldps
    killed t = S.fromList [ q | (ps, q) <- ldps, not (t `entails` ps) ]

    entails _ [] = True
    entails (t:ts) (p:ps) = (t `eq` p) && (ts `entails` ps)
    entails [] (_:_) = False

-- | Substitute state type by any enumerable type.
enumerateStates :: (Ord st, Ord a, Enum a) => FST st pred func -> FST a pred func
enumerateStates fst' =
  FST { fstS = S.fromList (M.elems statesMap)
      , fstE = edgesFromList
                 [ (aux q, lbl, aux q') | (q, lbl, q') <- edgesToList (fstE fst') ]
      , fstI = aux (fstI fst')
      , fstF = S.fromList [ aux q | q <- S.toList (fstF fst') ]
      }
  where
    statesMap = M.fromList (zip (S.toList (fstS fst')) [toEnum 0..])
    aux q = statesMap M.! q

accessibleStates :: (Ord q, PartialOrder pred, Boolean pred)
                 => OrderedEdgeSet q pred func -> S.Set q -> S.Set q
accessibleStates es initialWS = go initialWS S.empty
  where
    go ws acc
      | S.null ws = acc
      | (q, ws') <- S.deleteFindMin ws =
          let succs = S.fromList $ map snd (abstractEvalEdgesAll es q bot)
                                   ++ map snd (evalEpsilonEdges es q)
          in go (S.union ws' (S.difference succs acc)) (S.insert q acc)

coaccessibleStates :: (Ord q, PartialOrder pred, Boolean pred)
                   => OrderedEdgeSet q pred func -> S.Set q -> S.Set q
coaccessibleStates es = accessibleStates (flipEdges es)

trim :: (Ord q, PartialOrder pred, Boolean pred)
     => FST q pred func -> FST q pred func
trim fst' =
  FST { fstI = fstI fst'
      , fstF = S.intersection (fstF fst') useful
      , fstS = useful
      , fstE = edgesFromList [ (q, lbl, q') | (q, lbl, q') <- edgesToList (fstE fst')
                                            , S.member q useful
                                            , S.member q' useful ]
      }
  where
    useful = S.intersection (accessibleStates (fstE fst') (S.singleton (fstI fst')))
                            (coaccessibleStates (fstE fst') (fstF fst'))

-------------
-- Simulation
-------------

-- | Simulate a non-deterministic FST with single-symbol predicates
run :: (Function func
       ,SetLike pred (Dom func)
       ,Monoid (Rng func)
       ,Ord st)
       => FST st pred func -> [Dom func] -> [Rng func]
run fst' inp = do
  (os, q') <- go [ ([o], q') | (o,q') <- rightClosure fst' (fstI fst') ] inp
  guard (S.member q' (fstF fst'))
  return (mconcat $ reverse os)
  where
    go s [] = s
    go s (a:as) = go (close . step a $ s) as

    close s = prune S.empty [ (o':os, q') | (os, q) <- s, (o', q') <- rightClosure fst' q ]

    prune _   []        = []
    prune vis ((os,q):s) | S.member q vis = prune vis s
                         | otherwise      = (os,q):prune (S.insert q vis) s

    step a s = [ (o':os, q') | (os, q) <- s, (o',q') <- fstEvalEdges fst' q a ]

-- | Simulate a sequential machine with multi-symbol predicates
runSequential :: (Function func
                 ,Dom func ~ [a]
                 ,UniformListSet pred a
                 ,Monoid (Rng func)
                 ,Ord st)
              => FST st pred func -> [a] -> Maybe (Rng func)
runSequential fst' = go (fstI fst')
  where
    go q [] = if S.member q (fstF fst') then Just mempty else Nothing
    go q w | isChoiceState fst' q = do
               [(o, q')] <- pure (fstEvalEpsilonEdges fst' q)
               o' <- go q' w
               return $ mappend o o'
    go q w = do
      ts <- M.lookup q (eForward $ fstE fst')
      [(v, o, q')] <- pure [ (v, eval f u, q')
                           | (p, f, q') <- ts
                           , let n = listLength p
                           , let (u,v) = splitAt n w
                           , member u p ]
      o' <- go q' v
      return $ mappend o o'


runBacktracking :: (Monoid m, Function func, SetLike pred (Dom func), Stream s (Dom func))
       => (Rng func -> m)
       -> FST Int pred func
       -> P s m
runBacktracking phi fst' = lfp!(fstI fst') <* eof
  where
    lfp = M.fromList [ (q, go q) | q <- S.toList $ fstS fst' ]
    go q =
      (if isJoinState fst' q then barrier q else pure ())
      *> case fstEvalEpsilonEdges fst' q of
           [] -> case fstEdges fst' q of
                   [] -> pure mempty
                   es -> consume es
                 -- TODO: Only insert fBarrier if q is on an epsilon-loop
           es -> fBarrier q *> choose es

    consume es =
      foldl1 (<|>)
        [ (mappend . phi . eval func <$> litp (flip member p)) <*> (lfp!q')
        | (p, func, q') <- es ]

    choose [] = empty
    choose es = foldl1 (<|>) [ mappend (phi y) <$> (lfp!q') | (y, q') <- es ]

{-
-- | Example: This is how to interpret an action FST
interpAction :: (Function func
                ,SetLike pred Word8
                ,Dom func ~ Word8
                ,Rng func ~ [Either Word8 RegAction])
             => FST Int pred func -> ByteString -> Maybe Action
interpAction fst' b = runP' (runBacktracking (mconcat . map adjActionSem) fst') (0::Int, b)
-}
