import Data.List (intersect,union)

-- Definition of the types of Deterministic and 
-- Nondeterministic Finite Automata

-- Parameters: q type of states
--             sigma alphabet (type of symbols

-- Components: transition function, 
--             initial state(s), 
--             accepting states

data DFA q sigma = DFA (q -> sigma -> q) q [q]
data NFA q sigma = NFA (q -> sigma -> [q]) [q] [q]

deltaDFA :: DFA q sigma -> q -> sigma -> q
deltaDFA (DFA delta _ _) = delta

initialDFA :: DFA q sigma -> q
initialDFA (DFA _ q0 _) = q0

finalDFA :: DFA q sigma -> [q]
finalDFA (DFA _ _ fs) = fs

deltaNFA :: NFA q sigma -> q -> sigma -> [q]
deltaNFA (NFA delta _ _) = delta

initialNFA :: NFA q sigma -> [q]
initialNFA (NFA _ q0s _) = q0s

finalNFA :: NFA q sigma -> [q]
finalNFA (NFA _ _ fs) = fs


-- EXAMPLES

data Qex = Qex0 | Qex1 | Qex2 deriving (Eq,Enum,Bounded,Show)
data Sigma = A | B    deriving Show

dfa0 :: DFA Qex Sigma
dfa0 = DFA delta Qex0 [Qex0]
  where
    delta Qex0 A = Qex1
    delta Qex0 B = Qex0
    delta Qex1 A = Qex2
    delta Qex1 B = Qex1
    delta Qex2 A = Qex0
    delta Qex2 B = Qex2

nfa0 :: NFA Qex Sigma
nfa0 = NFA delta [Qex0] [Qex1,Qex2]
  where
    delta Qex0 A = [Qex0]
    delta Qex0 B = [Qex0,Qex1]
    delta Qex1 _ = [Qex2]
    delta Qex2 _ = []
    
-- Exercise 1: draw diagrams for the DFA dfa0 and the NFA nfa0

-- Exercise 2: define a DFA accepting words in which 
--   As and Bs have the same parity

data Qpar = Qpar0 | Qpar1 deriving (Eq,Enum,Bounded,Show)

dfa_parity :: DFA Qpar Sigma
dfa_parity = DFA delta Qpar0 [Qpar0]
  where
    delta Qpar0 _ = Qpar1
    delta Qpar1 _ = Qpar0

-- Exercise 3: define an NFA accepting words containing
--   the substring ABBA
 
data Qabba = Qeps | Qa | Qab | Qabb | Qabba deriving (Eq,Enum,Bounded,Show)
   
nfa_abba :: NFA Qabba Sigma
nfa_abba = NFA delta [Qeps] [Qabba]
  where
    delta Qeps A = [Qeps,Qa]
    delta Qeps B = [Qeps]
    delta Qa B = [Qab]
    delta Qab B = [Qabb]
    delta Qabb A = [Qabba]
    delta Qabba _ = [Qabba]
    delta _ _ = []

-- AUXILIARY FUNCTIONS

-- Represent a finite set as a list
-- Given a finite set, compute the set of its subsets
powerset :: [a] -> [[a]]
powerset [] = [[]]
powerset (a:as) = powerset as ++ map (a:) (powerset as)

-- Compute the list of all elements of a bounded type
--   We will use it to get the list of all states
allStates:: (Enum q, Bounded q) => [q]
allStates = [minBound..maxBound]

-- Union of all elements is a set of sets
-- union of two sets/lists is predefined
unions :: Eq a => [[a]] -> [a]
unions = foldl union []

-- Extended transition function for a DFA
-- (starDFA dfa q w) gives the state that we reach
--   after dfa processes the word w starting in state q

starDFA :: Eq q => DFA q sigma -> q -> [sigma] -> q
starDFA dfa q [] = q
starDFA dfa q (x:xs) = starDFA dfa (deltaDFA dfa q x) xs

-- Run a DFA on a word and determine if it accepts it or not

runDFA :: Eq q => DFA q sigma -> [sigma] -> Bool
runDFA dfa xs = starDFA dfa (initialDFA dfa) xs `elem` finalDFA dfa

-- Extended transition function for an NFA

starNFA :: Eq q => NFA q sigma -> [q] -> [sigma] -> [q]
starNFA nfa qs [] = qs
starNFA nfa qs (x:xs) = 
  starNFA nfa (unions [deltaNFA nfa q x | q <- qs]) xs

-- Accepting subsets of states of an NFA

acceptsNFA :: Eq q => NFA q sigma -> [q] -> Bool
acceptsNFA nfa qs = or (map (`elem` finalNFA nfa) qs)

-- Determine if a word is accepted by an NFA

runNFA :: Eq q => NFA q sigma -> [sigma] -> Bool
runNFA nfa xs = acceptsNFA nfa $ starNFA nfa (initialNFA nfa) xs

-- Trasform a DFA into an NFA (trivial)

dfa2nfa :: DFA q sigma -> NFA q sigma
dfa2nfa (DFA delta q0 fs) = NFA (\q x -> [delta q x]) [q0] fs

-- Transform an NFA into a DFA
--   The subset construction

nfa2dfa :: (Eq q, Enum q, Bounded q) => NFA q sigma -> DFA [q] sigma
nfa2dfa nfa = DFA deltaD 
                  (initialNFA nfa) 
                  (filter (acceptsNFA nfa) (powerset allStates))
  where
    deltaD qs x = unions [deltaNFA nfa q x | q <- qs]