// NFAs and DFAs based on Scala's partial functions


// (1) Write a polymorphic function that tests whether the 
// intersection of two sets is non-empty

def share[A](a: Set[A], b: Set[A]) : Boolean =
  !(a intersect b).isEmpty

share(Set(1,2,3), Set(2, 3, 4))   // true
share(Set(1,2,3), Set(4, 5, 6))   // false


// State nodes of the DFAs and NFAs
abstract class State
type States = Set[State]

// Some states for test cases
case object Q0 extends State
case object Q1 extends State
case object Q2 extends State
case object Q3 extends State
case object Q4 extends State
case object Q5 extends State
case object Q6 extends State


// Transitions for DFAs and NFAs
type Trans = PartialFunction[(State, Char), State]
type NTrans = Set[Trans]


// example transition of an DFA
val dtrans : Trans = 
  { case (Q0, 'a') => Q1 
    case (Q0, 'b') => Q0
    case (Q1, 'a') => Q2 
    case (Q1, 'b') => Q0
    case (Q2, 'a') => Q2 
    case (Q2, 'b') => Q0 
  }


// (2) Write a function that takes a transition and a 
// (state, character)-pair as arguments and produces an 
// optional state (the state specified by the partial transition
// function whenever it is defined; if the transition function 
// is undefined, return None.

def fire(e: Trans, qc: (State, Char)) : Option[State] = 
  e.lift.apply(qc)


// (3) Write a function that takes a transition, a state 
// and a list of characters as arguments and produces 
// the state generated by following the transitions for 
// each character in the list.

def nexts(trans: Trans, q: State, s: List[Char]) : Option[State] = s match {
  case Nil => Some(q)
  case c::cs => fire(trans, (q, c)).flatMap(nexts(trans, _, cs))
}



// class for DFAs
case class DFA(start: State,      // starting state
               trans: Trans,      // transition
               fins:  States)     // final states

// (4) Write a function that tests whether a string is accepted 
// by an DFA or not.

def accepts(dfa: DFA, s: String) : Boolean = nexts(dfa.trans, dfa.start, s.toList) match {
  case None => false
  case Some(q) => dfa.fins contains q
}


// DFA examples
 
val dtrans1 : Trans = 
  { case (Q0, 'a') => Q0 
    case (Q0, 'b') => Q1 
  }

val dfa1 = DFA(Q0, dtrans1, Set[State](Q1))

accepts(dfa1, "aaab")     // true
accepts(dfa1, "aacb")     // false


// NFAs


// (5) Write a function that takes a transition set, a state
// and a character as arguments, and calculates all possible 
// next states (returned as set).

def nnext(trans: NTrans, q: State, c: Char) : States = {
  trans.map(fire(_, (q, c))).flatten
}

// (6) Write a function that takes a transition set, a set of states 
// and a character as arguments, and calculates all possible 
// next states that can be reached from any state in the set.

def nnexts(trans: NTrans, qs: States, c: Char) : States = {
  qs.flatMap(nnext(trans, _, c))
}


// (7) Write a function that lifts nnexts from from single 
// characters to lists of characters.
def nnextss(trans: NTrans, qs: States, s: List[Char]) : States = s match {
  case Nil => qs
  case c::cs => {
    val ns = nnexts(trans, qs, c)
    nnextss(trans, ns, cs) 
  }
}

// class for NFAs
case class NFA(start: States,       // starting state
               trans: NTrans,       // transition edges
               fins:  States)       // final states


// (8) Write a function that tests whether a string is 
// accepted by an NFA or not.

def naccepts(nfa: NFA, s: String) : Boolean = {
  share(nnextss(nfa.trans, nfa.start, s.toList), nfa.fins)
}


// (9) Write similar functions as in (7) and (8), but instead of 
// returning states or a boolean, calculate the number of states 
// that need to be followed in each step.

def max_nextss(trans: NTrans, qs: States, s: List[Char], max: Int) : Int = s match {
  case Nil => max
  case c::cs => {
    val ns = nnexts(trans, qs, c)
    val ns_size = ns.size
    if (max < ns_size) max_nextss(trans, ns, cs, ns_size) 
    else max_nextss(trans, ns, cs, max)
  }
}

def max_accepts(nfa: NFA, s: String) : Int = {
  max_nextss(nfa.trans, nfa.start, s.toList, 0)
}


// NFA examples


// 1 
val trans1 : NTrans = Set(
  { case (Q0, 'a') => Q1 },
  { case (Q0, _)   => Q0 },
  { case (Q1, _)   => Q2 },
  { case (Q2, _)   => Q3 },
  { case (Q3, _)   => Q4 },
  { case (Q4, 'b') => Q5 },
  { case (Q5, 'c') => Q6 }
)

val nfa1 = NFA(Set[State](Q0), trans1, Set[State](Q6))

naccepts(nfa1, "axaybzbc")     // true
naccepts(nfa1, "aaaaxaybzbc")  // true
naccepts(nfa1, "axaybzbd")     // false

// the nfa has five states, which might be all 
// active

max_accepts(nfa1, "axaybzbc")     // 3 
max_accepts(nfa1, "aaaaxaybzbc")  // 5
max_accepts(nfa1, "axaybzbd")     // 3
max_accepts(nfa1, "aaaaaaaaaaaaaxaybzbd")   // 5


// 2
val trans2 : NTrans = Set(
  { case (Q0, 'a') => Q0 },
  { case (Q0, 'a') => Q1 },
  { case (Q0, 'b') => Q2 },
  { case (Q1, 'a') => Q1 },
  { case (Q2, 'b') => Q2 }
)

val nfa2 = NFA(Set[State](Q0), trans2, Set[State](Q2))

naccepts(nfa2, "aa")             // false
naccepts(nfa2, "aaaaa")          // false
naccepts(nfa2, "aaaaab")         // true
naccepts(nfa2, "aaaaabbb")       // true
naccepts(nfa2, "aaaaabbbaaa")    // false
naccepts(nfa2, "ac")             // false

// 3
val trans3 : NTrans = Set(
  { case (Q0, _)   => Q0 },
  { case (Q0, 'a') => Q1 },
  { case (Q0, 'b') => Q3 },
  { case (Q1, 'b') => Q2 },
  { case (Q2, 'c') => Q5 },
  { case (Q3, 'c') => Q4 },
  { case (Q4, 'd') => Q5 }
)

val nfa3 = NFA(Set[State](Q0), trans3, Set[State](Q5))

naccepts(nfa3, "aaaaabc")      // true
naccepts(nfa3, "aaaabcd")      // true
naccepts(nfa3, "aaaaab")       // false
naccepts(nfa3, "aaaabc")       // true
naccepts(nfa3, "aaaaabbbaaa")  // false