{-# LANGUAGE ScopedTypeVariables, MonadComprehensions, DeriveAnyClass,
             DeriveGeneric #-}
{-|
  Module      : Text.ANTLR.Lex.NFA
  Description : Nondeterministic finite automatons and algorithms to compute DFAs
  Copyright   : (c) Karl Cronburg, 2018
  License     : BSD3
  Maintainer  : karl@cs.tufts.edu
  Stability   : experimental
  Portability : POSIX

-}
module Text.ANTLR.Lex.NFA where
import Text.ANTLR.Lex.Automata
import Text.ANTLR.Lex.DFA (DFA(..))
import qualified Text.ANTLR.Lex.DFA as DFA

import Text.ANTLR.Set (singleton, notMember, union, Set(..), member, Hashable)
import qualified Text.ANTLR.Set as Set
import Text.ANTLR.Set (fromList, toList)

import Data.List (maximumBy)
import GHC.Generics (Generic)

-- | NFA edges can be labeled with either a symbol in symbol alphabet @s@,
--   or an epsilon.
data Edge s = Edge s | NFAEpsilon
  deriving (Ord, Eq, Hashable, Generic)

instance (Show s) => Show (Edge s) where
  show NFAEpsilon = "ϵ"
  show (Edge s)   = "E(" ++ show s ++ ")"

-- | Is this an edge (not an epsilon)?
isEdge :: Edge s -> Bool
isEdge (Edge _) = True
isEdge _        = False

-- | An NFA is an automata with edges @'Edge' s@ and nodes @i@.
type NFA s i = Automata (Edge s) s i

-- | NFA states
type State i = i

-- | DFA states as constructed from an NFA is a set (config) of NFA states.
type DFAState i = Config (State i)

-- | Epsilon closure of an NFA is a closure where we can traverse epsilons.
epsClosure ::
  (Ord i, Hashable i, Hashable s, Eq s)
  => Automata (Edge s) s i -> Config i -> Config i
epsClosure = closureWith (NFAEpsilon ==)

-- | Subset construction algorithm for constructing a DFA from an NFA.
nfa2dfa_slow :: forall s i. (Hashable s, Eq s, Hashable i, Eq i, Ord i)
  => NFA s i -> DFA s (Set (State i))
nfa2dfa_slow nfa@Automata{s0 = s0, _Σ = _Σ, _F = _F0} = let
    
    epsCl = epsClosure nfa
    mv    = move nfa

    dS :: Config (DFAState i) -> Config (DFAState i) -> Set (Transition (DFA.Edge s) (DFAState i))
    dS marked ts
      | Set.null ts = Set.empty
      | otherwise = let
        
          _Δ  = fromList
                [ (_T, (False, singleton a), epsCl (mv _T (Edge a)))
                | _T <- toList ts
                , _T `notMember` marked
                , a  <- toList _Σ
                ]

          _Us = Set.map (\(a,b,c) -> c) _Δ
          fromStates = Set.map (\(a,b,c) -> a) _Δ

        in _Δ `union` dS (fromStates `union` marked) _Us
    
    _Δ' :: Set (Transition (DFA.Edge s) (DFAState i))
    _Δ' = dS Set.empty (singleton s0')

    s0' = epsCl $ singleton s0

  in Automata
      { _S = fromList [ tFrom x | x <- toList _Δ' ] `union` fromList [ tTo x | x <- toList _Δ' ]
      , _Σ = _Σ
      , _Δ = _Δ'
      , s0 = s0'
      , _F = fromList [nfaState | (_,_,nfaState) <- toList _Δ', c <- toList nfaState, c `member` _F0]
      }

-- | Subset construction but where we compress our sets of transitions along the way.
nfa2dfa :: forall s i. (Hashable s, Eq s, Hashable i, Eq i, Ord i)
  => NFA s i -> DFA s (Set (State i))
nfa2dfa nfa@Automata{s0 = s0, _Σ = _Σ, _S = _S, _F = _F0} = let
    
    epsCl = epsClosure nfa
    mv    = move nfa

    dS :: Config (DFAState i) -> Config (DFAState i) -> Set (Transition (DFA.Edge s) (DFAState i))
    dS marked ts
      | Set.null ts = Set.empty
      | otherwise = let
        
          _Δ  =
            Set.fromList
                [ (_T, (False, singleton a), epsCl (mv _T (Edge a)))
                | _T <- Set.toList ts
                , _T `notMember` marked
                , a  <- Set.toList _Σ
                ]
            `union`
            Set.fromList
                [ (_T, (True, _Σ), epsCl $ moveComplement nfa _T)
                | _T <- Set.toList ts
                , _T `notMember` marked
                ]

          _Us = fromList [ c | (a,b,c) <- toList _Δ ]
          fromStates = fromList [ a | (a,b,c) <- toList _Δ ]

        in _Δ `union` dS (fromStates `union` marked) _Us
    
    _Δ' :: Set (Transition (DFA.Edge s) (DFAState i))
    _Δ' = let run_dS = dS Set.empty (singleton s0')
          in  Set.filter (\(_, _, b) -> not $ Set.null b) $ compress run_dS

    s0' = epsCl $ singleton s0

  in Automata
      { _S = fromList [ tFrom x | x <- toList _Δ' ] `union` fromList [ tTo x | x <- toList _Δ' ]
      , _Σ = _Σ
      , _Δ = _Δ'
      , s0 = s0'
      , _F = fromList [nfaState | (_,_,nfaState) <- toList _Δ', c <- toList nfaState, c `member` _F0]
      }

-- | Compute all the states statically used in a particular set of transitions.
allStates :: forall s i. (Hashable i, Eq i) => Set (Transition (Edge s) i) -> Set (State i)
allStates ts = fromList [ n | (n, _, _) <- toList ts ] `union` fromList [ n | (_, _, n) <- toList ts ]

-- | Converts the given list of transitions into a complete NFA / Automata
--   structure, assuming two things:
--
-- > The first node of the first edge is the start state
-- > The last  node of the last  edge is the (only) final state
--
list2nfa :: forall s i. (Hashable i, Eq i, Hashable s, Eq s) => [Transition (Edge s) i] -> NFA s i
list2nfa [] = undefined
list2nfa ((t@(n1,_,_)):ts) = Automata
  { _S = allStates $ Set.fromList (t:ts)
  , _Σ = Set.fromList [ e
          | (_, es, _) <- t:ts
          , Edge e     <- filter isEdge (Set.toList $ snd es)
          ]
  , s0 = n1
  , _F = Set.fromList [ (\(_,_,c) -> c) $ last (t:ts) ]
  , _Δ = Set.fromList $ t:ts
  }

-- | Rename the states in the second NFA such that they start at the index
--   one greater than the maximum index of the first NFA.
shiftAllStates ::
  forall s i. (Hashable i, Eq i, Ord i, Hashable s, Eq s)
  => (i -> Int) -> (Int -> i) -> NFA s i -> NFA s i -> NFA s i
shiftAllStates from to
  n1 (n2@Automata{_Δ = _Δ2, _S = _S2, _F = _F2, s0 = s2_0})
  = n2 { _Δ = fromList [ (to $ from i0 + shift, e, to $ from i1 + shift) | (i0, e, i1) <- toList _Δ2 ]
       , _S = fromList [ to $ from i + shift | i <- toList _S2 ]
       , _F = fromList [ to $ from i + shift | i <- toList _F2 ]
       , s0 = to $ from s2_0 + shift
       }
  where
    shift = 1 + foldr (\(i0, _, i1) i -> from $ maximum [to i, i0, i1]) 0 (_Δ n1)

-- | Take the union of two NFAs, renaming states according to 'shiftAllStates'.
nfaUnion ::
  forall s i. (Ord i, Hashable i, Eq i, Hashable s, Eq s)
  => (i -> Int) -> (Int -> i) -> NFA s i -> NFA s i -> NFA s i
nfaUnion from to
  (n1@Automata{_Δ = _Δ1, _S = _S1, _F = _F1, s0 = s1_0}) n2
  = let

    Automata{_Δ = _Δ2, _S = _S2, _F = _F2, s0 = s2_0} = shiftAllStates from to n1 n2
    mx2 = 1 + foldr (\(i0, _, i1) i -> from $ maximum [to i, i0, i1]) 0 _Δ2

    _Δ' =     _Δ1
      `union` _Δ2
      `union` Set.singleton (s0', (False, singleton NFAEpsilon), s1_0)
      `union` Set.singleton (s0', (False, singleton NFAEpsilon), s2_0)
      `union` fromList [ (f1_0, (False, singleton NFAEpsilon), f0') | f1_0 <- toList _F1 ]
      `union` fromList [ (f2_0, (False, singleton NFAEpsilon), f0') | f2_0 <- toList _F2 ]

    s0' = to mx2
    f0' = to $ mx2 + 1

  in Automata
    { _S = allStates _Δ'
    , _Σ =  fromList [ e
            | (_, es, _)  <- toList _Δ'
            , Edge e      <- toList $ Set.filter isEdge $ snd es
            ]
    , s0 = s0'
    , _F = Set.fromList [f0']
    , _Δ = _Δ'
    }

-- | Concatenate two NFAs, renaming states in the second NFA according to 'shiftAllStates'.
nfaConcat ::
  forall s i. (Hashable i, Eq i, Ord i, Hashable s, Eq s) => (i -> Int) -> (Int -> i) -> NFA s i -> NFA s i -> NFA s i
nfaConcat from to
  (n1@Automata{_Δ = _Δ1, _S = _S1, _F = _F1, s0 = s1_0}) n2
  = let
    Automata{_Δ = _Δ2, _S = _S2, _F = _F2, s0 = s2_0} = shiftAllStates from to n1 n2
    
    _Δ' =     _Δ1
      `union` _Δ2
      `union` fromList [ (f1_0, (False, singleton NFAEpsilon), s2_0) | f1_0 <- toList _F1 ]
  
  in Automata
    { _S = allStates _Δ'
    , _Σ =  fromList [ e
            | (_, es, _)  <- toList _Δ'
            , Edge e      <- toList $ Set.filter isEdge $ snd es
            ]
    , s0 = s1_0
    , _F = _F2
    , _Δ = _Δ'
    }

-- | Take the Kleene-star of an NFA, adding epsilons as needed.
nfaKleene :: forall s i. (Ord i, Hashable i, Eq i, Hashable s, Eq s) => (i -> Int) -> (Int -> i) -> NFA s i -> NFA s i
nfaKleene from to 
  (n1@Automata{_Δ = _Δ1, _S = _S1, _F = _F1, s0 = s1_0})
  = let
    mx1 = 1 + foldr (\(i0, _, i1) i -> from $ maximum [to i, i0, i1]) 0 _Δ1

    s0' = to mx1
    f0' = to $ mx1 + 1

    _Δ' =     _Δ1
      `union` Set.singleton (s0', (False, singleton NFAEpsilon), s1_0)
      `union` Set.singleton (s0', (False, singleton NFAEpsilon), f0')
      `union` fromList [ (f1_0, (False, singleton NFAEpsilon), s1_0) | f1_0 <- toList _F1 ]
      `union` fromList [ (f1_0, (False, singleton NFAEpsilon), f0')  | f1_0 <- toList _F1 ]

  in Automata
    { _S = allStates _Δ'
    , _Σ =  fromList [ e
            | (_, es, _)  <- toList _Δ'
            , Edge e      <- toList $ Set.filter isEdge $ snd es
            ]
    , s0 = s0'
    , _F = Set.fromList [f0']
    , _Δ = _Δ'
    }