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

{-|
Description: Finite state transducers with epsilon-transitions
-}
module Control.FST
  ( FST(..)
  -- * Constructing FSTs
  , wordFST
  -- * Operations on FSTs
  , productFST, invertFST, restrictFST
  -- * Debugging Utilities
  , liveFST
  ) 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.Maybe (mapMaybe, fromMaybe)

import Numeric.Natural

import Control.Lens.TH

import Control.Monad.State.Strict

import Text.PrettyPrint.Leijen (Pretty(..))
import qualified Text.PrettyPrint.Leijen as PP

import Debug.Trace

data FST state input output = FST
  { stInitial :: Set state
  , stTransition :: Map (state, Maybe input) (Set (state, Maybe output))
  , stAccept :: Set state
  } deriving (Show, Read)

instance (Show state, Show input, Show output) => Pretty (FST state input output) where
  pretty FST{..} = PP.vsep
    [ PP.text "Initial states:" PP.</> PP.hang 2 (PP.list . map (PP.text . show) $ Set.toAscList stInitial)
    , PP.text "State transitions:" PP.<$> PP.indent 2 (PP.vsep
        [        PP.text (show st)
          PP.<+> (PP.text "-" PP.<> PP.tupled [label inS, label outS] PP.<> PP.text "→")
          PP.<+> PP.text (show st')
        | ((st, inS), to) <- Map.toList stTransition
        , (st', outS) <- Set.toAscList to
        ])
    , PP.text "Accepting states:" PP.</> PP.hang 2 (PP.list . map (PP.text . show) $ Set.toAscList stAccept)
    ]
    where
      label :: Show a => Maybe a -> PP.Doc
      label = maybe (PP.text "ɛ") (PP.text . show)

wordFST :: forall input output. Seq output -> FST Natural input output
-- ^ @wordFST str@ is the minimal FST generating @str@ as output when given no input
wordFST outs = FST
  { stInitial    = Set.singleton 0
  , stAccept     = Set.singleton l
  , stTransition = Map.fromSet next states
  }
  where
    l :: Natural
    l = fromIntegral $ Seq.length outs
    states :: Set (Natural, Maybe input)
    states = Set.fromDistinctAscList [ (n, Nothing) | n <- [0..pred l] ]
    next :: (Natural, Maybe input) -> Set (Natural, Maybe output)
    next (i, _) = Set.singleton (succ i, Just . Seq.index outs $ fromIntegral i)

productFST :: forall state1 state2 input output. (Ord state1, Ord state2, Ord input, Ord output) => FST state1 input output -> FST state2 input output -> FST (state1, state2) input output
-- ^ Cartesian product on states, logical conjunction on transitions and state-properties (initial and accept)
--
-- This is most intuitive when thought of as the component-wise product of weighted FSTs with weights in the boolean semiring.
productFST fst1 fst2 = FST
  { stInitial    = stInitial fst1 `setProduct` stInitial fst2
  , stAccept     = stAccept fst1 `setProduct` stAccept fst2
  , stTransition = Map.fromSet transitions . Set.fromList . mapMaybe filterTransition . Set.toAscList $ Map.keysSet (stTransition fst1) `setProduct` Map.keysSet (stTransition fst2)
  }
  where
    setProduct :: forall a b. Set a -> Set b -> Set (a, b)
    setProduct as bs = Set.fromDistinctAscList $ (,) <$> Set.toAscList as <*> Set.toAscList bs
    filterTransition :: forall label. Eq label => ((state1, Maybe label), (state2, Maybe label)) -> Maybe ((state1, state2), Maybe label)
    filterTransition ((st1, Nothing ), (st2, in2     )) = Just ((st1, st2), in2)
    filterTransition ((st1, in1     ), (st2, Nothing )) = Just ((st1, st2), in1)
    filterTransition ((st1, Just in1), (st2, Just in2))
      | in1 == in2 = Just ((st1, st2), Just in1)
      | otherwise  = Nothing
    transitions :: ((state1, state2), Maybe input) -> Set ((state1, state2), Maybe output)
    transitions ((st1, st2), inS) = Set.fromList . mapMaybe filterTransition . Set.toAscList $ out1 `setProduct` out2
      where
        out1 = (fromMaybe Set.empty $ stTransition fst1 !? (st1, inS)) `Set.union` (fromMaybe Set.empty $ stTransition fst1 !? (st1, Nothing))
        out2 = (fromMaybe Set.empty $ stTransition fst2 !? (st2, inS)) `Set.union` (fromMaybe Set.empty $ stTransition fst2 !? (st2, Nothing))

restrictFST :: forall state input output. (Ord state, Ord input, Ord output) => Set state -> FST state input output -> FST state input output
restrictFST sts FST{..} = FST
  { stInitial    = stInitial `Set.intersection` sts
  , stAccept     = stAccept  `Set.intersection` sts
  , stTransition = Map.mapMaybeWithKey restrictTransition stTransition
  }
  where
    restrictTransition :: (state, Maybe input) -> Set (state, Maybe output) -> Maybe (Set (state, Maybe output))
    restrictTransition (st, _) tos = tos' <$ guard (st `Set.member` sts)
      where
        tos' = Set.filter (\(st', _) -> st' `Set.member` sts) tos
      

liveFST :: forall state input output. (Ord state, Ord input, Ord output, Show state) => FST state input output -> Set state
-- ^ Compute the set of "live" states (with no particular complexity)
--
-- A state is "live" iff there is a path from it to an accepting state and a path from an initial state to it
liveFST fst@FST{..} = flip execState Set.empty $ mapM_ (depthSearch Set.empty) stInitial
  where
    stTransition' :: Map state (Set state)
    stTransition' = Map.map (Set.map $ \(st, _) -> st) $ Map.mapKeysWith Set.union (\(st, _) -> st) stTransition
    depthSearch :: Set state -> state -> State (Set state) ()
    depthSearch acc curr = do
      let acc' = Set.insert curr acc
          next = fromMaybe Set.empty $ stTransition' !? curr
      alreadyLive <- get
      when (not . Set.null $ Set.insert curr next `Set.intersection` Set.union stAccept alreadyLive) $
        modify $ Set.union acc'
      mapM_ (depthSearch acc') $ next `Set.difference` alreadyLive


invertFST :: FST state input output -> Seq output -> Set (Seq input)
invertFST = undefined
\end{code}