module Automaton.NFAe (
    NFAe(..)
  , isAcceptState
  , step'
  , ecl
  , alphabets
  , states
  , mapState
  , module Automaton
  ) where

import Data.Maybe (catMaybes)

import qualified Data.Map as M
import qualified Data.Set as S

import Automaton

-- ε遷移つき非決定性有限オートマトン
data NFAe state alpha = NFAe {
  -- Nothing が空列εを表す
    nfaeTransTable :: M.Map (state, Maybe alpha) (S.Set state)
  , nfaeStart :: state
  , nfaeAccepts :: S.Set state
  } deriving (Show, Read)

isAcceptState :: (State state, Alphabet alpha)
                 => NFAe state alpha -> state -> Bool
isAcceptState nfae s = S.member s $ nfaeAccepts nfae

-- デフォルト付き lookup
lookup' :: Ord k => k -> M.Map k (S.Set a) -> S.Set a
lookup' k map = case M.lookup k map of
  Nothing -> S.empty
  Just s -> s

-- 状態 s で文字 a を入力したときの(ε遷移も含めた)遷移先を返す。
step' :: (State state, Alphabet alpha)
         => NFAe state alpha -> state -> alpha -> S.Set state
step' nfae s a = ecl nfae $ lookup' (s, Just a) $ nfaeTransTable nfae

-- ε閉包 (ε-cl)
-- 状態集合 states 自身と、そこからε遷移により辿り着ける全ての状態を返す。
ecl :: (State state, Alphabet alpha)
              => NFAe state alpha -> S.Set state -> S.Set state
ecl nfae@NFAe{..} states
  | states == states' = states
  | otherwise = ecl nfae states'
  where
    states' = S.fold iter states states
    -- s からε遷移する状態を付け加える。
    iter s states = S.union states $ lookup' (s, Nothing) nfaeTransTable

instance Automaton NFAe where
  type F NFAe s = S.Set s
  step nfae states a =
    S.fold (\s states' -> S.union states' $ step' nfae s a) S.empty states
  initStates nfae = S.singleton $ nfaeStart nfae
  isAcceptStates nfae states = any (isAcceptState nfae) $ S.toList states

-- nfae が認識する言語の文字集合
alphabets :: (State state, Alphabet alpha)
             => NFAe state alpha -> S.Set alpha
alphabets nfae =
  S.fromList $ catMaybes $ map snd $ M.keys $ nfaeTransTable nfae

-- nfae の状態の集合
states :: (State state, Alphabet alpha)
          => NFAe state alpha -> S.Set state
states nfae = S.union doms rans
  where
    -- doms : 遷移関数の定義域, rans : 値域
    doms = S.fromList $ map fst $ M.keys $ nfaeTransTable nfae
    rans = foldl S.union S.empty $ M.elems $ nfaeTransTable nfae

-- NFAe の状態を map する。ただし整合性は自分で保証してね。
mapState :: (State state, State state', Alphabet alpha)
            => (state -> state') -> NFAe state alpha -> NFAe state' alpha
mapState f NFAe{..} = NFAe {
    nfaeTransTable =
        M.mapKeys (\(s, a) -> (f s, a)) $ M.map (S.map f) nfaeTransTable
  , nfaeStart = f nfaeStart
  , nfaeAccepts = S.map f nfaeAccepts
  }
