import scala.language.implicitConversions

import scala.language.reflectiveCalls

abstract class Rexp

case object NULL extends Rexp

case object EMPTY extends Rexp

case class CHAR(c: Char) extends Rexp

case class ALT(r1: Rexp, r2: Rexp) extends Rexp 

case class SEQ(r1: Rexp, r2: Rexp) extends Rexp 

case class STAR(r: Rexp) extends Rexp 

abstract class Pat

case class VAR(x: String, r: Rexp) extends Pat

case class GRP(x: String, p: Pat) extends Pat

case class PSEQ(p1: Pat, p2: Pat) extends Pat

case class PALT(p1: Pat, p2: Pat) extends Pat

case class PSTAR(p: Pat) extends Pat

def nullable (r: Rexp) : Boolean = r match {

  case NULL => false

  case EMPTY => true

  case CHAR(_) => false

  case ALT(r1, r2) => nullable(r1) || nullable(r2)

  case SEQ(r1, r2) => nullable(r1) && nullable(r2)

  case STAR(_) => true

}

def down (p: Pat) : Rexp = p match {

  case VAR(x: String, w: String, r: Rexp) => r

  case GRP(x: String, w: String, p: Pat) => down(p)

  case PSEQ(p1: Pat, p2: Pat) => SEQ(down(p1), down(p2))

  case PALT(p1: Pat, p2: Pat) => ALT(down(p1), down(p2))

  case PSTAR(p: Pat) => STAR(down(p))

}

def patnullable (p: Pat) : Boolean = p match {

  case PVar(_, r) => nullable(r)

  case PSEQ(p1, p2) => patnullable(p1) && patnullable(p2)

  case PCHOICE(p1, p2) => patnullable(p1) || patnullable(p2)

  case PSTAR(p) => true

  case PatVar(_, p) => patnullable(p)

}

//type Env = Set[List[(String, String)]]

type Env = Map[Int, String]

def update(n: Int, c: Char) (env: Env) = 

  env + (n -> (env.getOrElse(n, "") + c.toString))

def pderivPat (p: Pat, c: Char) : Set[(Pat, Env => Env)] = p match {

  case PVar(n: Int, r: Rexp) => {

    val pds = pderiv(r, c)

    if (pds.isEmpty) Set()

    else Set((PVar(n, toRexp(pds.toList)), update(n, c)))

  }

  case PSEQ(p1: Pat, p2: Pat) => {

    val pats : Set[(Pat, Env => Env)] = 

      for ((p, f) <- pderivPat(p1, c)) yield (PSEQ(p, p2), f)

    if (nullable(strip(p1))) pats ++ pderivPat(p2, c)  

    else pats

  }

  case PCHOICE(p1: Pat, p2: Pat) => pderivPat(p1, c) ++ pderivPat(p2, c)

  case PSTAR(p1: Pat) => 

    for ((p, f) <- pderivPat(p1, c)) yield (PSEQ(p, PSTAR(p1)), f)

  case PatVar(n: Int, p1: Pat) => 

    for ((p, f) <- pderivPat(p1, c)) yield (PatVar(n, p), f compose (update (n, c)))

}

val p2 = PSEQ(PVar(1, STAR("A")), PVar(2, STAR("A")))

pderivPat(p2, 'A').mkString("\n")

def greedy_aux(env: Set[(Pat, Env)], w: List[Char]) : Set[Env] = w match {

  case Nil => 

    for ((p, e) <- env if patnullable(p)) yield e

  case c::cs => {

    val env2 = for ((p, e) <- env; (p1, f) <- pderivPat(p, c)) yield (p1, f(e))

    greedy_aux(env2, cs)

  }

}

def greedy(p: Pat, w: String) = {

  val res = greedy_aux(Set((p, Map())), w.toList)

  if (res.isEmpty) None else Some(res)

}

// some convenience for typing in regular expressions

def charlist2rexp (s : List[Char]) : Rexp = s match {

  case Nil => EMPTY

  case c::Nil => CHAR(c)

  case c::s => SEQ(CHAR(c), charlist2rexp(s))

}

implicit def string2rexp (s : String) : Rexp = charlist2rexp(s.toList)

implicit def RexpOps (r: Rexp) = new {

  def | (s: Rexp) = ALT(r, s)

  def % = STAR(r)

  def ~ (s: Rexp) = SEQ(r, s)

}

implicit def stringOps (s: String) = new {

  def | (r: Rexp) = ALT(s, r)

  def | (r: String) = ALT(s, r)

  def % = STAR(s)

  def ~ (r: Rexp) = SEQ(s, r)

  def ~ (r: String) = SEQ(s, r)

}

implicit def PatOps (p: Pat) = new {

  def | (q: Pat) = PALT(p, q)

  def % = PSTAR(p)

  def ~ (q: Pat) = PSEQ(p, q)

}

val p3 = PSTAR(PSEQ(PVar(1, "A"), PVar(2, "B")))

greedy2(Set((p3, Map())), "ABAB".toList)

val p4 = PVar(1, "A")

greedy2(Set((p4, Map())), "A".toList)

val p5 = PSEQ(PVar(1, "A"), PVar(1, "B"))

greedy2(Set((p5, Map())), "AB".toList)

val res = pderivPat(p5, 'A')

for ((_, f) <- res) yield f(Map())

val p6 = PatVar(4,PSTAR(PCHOICE(PCHOICE(PVar(1, "A"), PVar(2, "AB")), PVar(3, "B"))))

greedy(p6, "ABA")