// 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