\begin{code}
{-# LANGUAGE ScopedTypeVariables
#-}

{-|
Description: Deterministic finite state transducers
-}
module Control.DFST
  ( DFST(..)
  , runDFST, runDFST'
  , toFST
  ) where

import Data.Map.Strict (Map, (!?))
import qualified Data.Map.Strict as Map

import Data.Set (Set)
import qualified Data.Set as Set

import Data.Sequence (Seq(..))
import qualified Data.Sequence as Seq

import Data.Monoid

import Numeric.Natural

import Control.Monad
import Control.Monad.State

import Control.FST (FST(FST))
import qualified Control.FST as FST


data DFST state input output = DFST
  { stInitial :: state
  , stTransition :: Map (state, input) (state, Seq output)
    -- ^ All @(s, c)@-combinations not mapped are assumed to map to @(s, Nothing)@
  , stAccept :: Set state
  }


toFST :: forall state input output. (Ord state, Ord input, Ord output) => DFST state input output -> FST (state, Maybe (input, Natural)) input output
-- ^ Split apart non-singleton outputs into a series of epsilon-transitions
toFST DFST{..} = flip execState initialFST $ forM_ (Map.toList stTransition) handleTransition
  where
    initialFST = FST
      { stInitial = Set.singleton (stInitial, Nothing)
      , stTransition = Map.empty
      , stAccept = Set.map (, Nothing) stAccept
      }
    addTransition :: forall state input output. (Ord state, Ord input, Ord output) => (state, Maybe input) -> (state, Maybe output) -> State (FST state input output) ()
    addTransition k v = modify $ \f@FST{ stTransition } -> f { FST.stTransition = Map.insertWith Set.union k (Set.singleton v) stTransition }
    handleTransition :: ((state, input), (state, Seq output)) -> State (FST (state, Maybe (input, Natural)) input output) ()
    handleTransition ((st, inS), (st', outs)) = handleTransition' (st, Nothing) (Just inS) outs (st', Nothing)
    handleTransition' :: (state, Maybe (input, Natural)) -> Maybe input -> Seq output -> (state, Maybe (input, Natural)) -> State (FST (state, Maybe (input, Natural)) input output) ()
    handleTransition' from inS Empty to = addTransition (from, inS) (to, Nothing)
    handleTransition' from inS (outS :<| Empty) to = addTransition (from, inS) (to, Just outS)
    handleTransition' from@(st, chain) inS (outS :<| oo) to = do
      let next
            | Just (inS', i) <- chain = (st, Just (inS', succ i))
            | Just inS'      <- inS   = (st, Just (inS', 0     ))
              -- Both calls to `handleTransition'` (one in `handleTransition`, the other below) satisfy one of the above cases
      addTransition (from, inS) (next, Just outS)
      handleTransition' next Nothing oo to
  
runDFST :: forall state input output. (Ord state, Ord input) => DFST state input output -> Seq input -> Maybe (Seq output)
runDFST dfst@DFST{..} str = let (finalState, str') = runDFST' dfst stInitial str Seq.empty
                            in str' <$ guard (finalState `Set.member` stAccept)

runDFST' :: forall state input output. (Ord state, Ord input)
         => DFST state input output
         -> state               -- ^ Current state
         -> Seq input           -- ^ Remaining input
         -> Seq output          -- ^ Accumulator containing previous output
         -> (state, Seq output) -- ^ Next state, altered output
runDFST' _             st  Empty     acc = (st, acc)
runDFST' dfst@DFST{..} st (c :<| cs) acc
  | Just (st', mc') <- stTransition !? (st, c)
  = runDFST' dfst st' cs $ acc <> mc'
  | otherwise
  = runDFST' dfst st cs acc
\end{code}