import TransitionSystem._
import Expr._
import Utils._
import ScalaZ3Interface._

import sys.process._
import sys.error

import scala.language.postfixOps

case class TransitionSystem(
  val stateIds: Seq[StateIdentifier],
  val auxiliaryIds: Seq[AuxiliaryIdentifier],
  val inputIds: Seq[InputIdentifier],
  val transition: Map[Identifier, Expr],
) {

  checkWF()

  override def toString: String = {
    (auxiliaryIds ++ stateIds).map(id =>
      s"$id = ${transition(id)}"
    ).mkString("\n")
  }

  // Check some well-formedness conditions on the transition system
  def checkWF(): Unit = {
    // No two variables should have the same name
    val allVars = (stateIds ++ auxiliaryIds ++ inputIds).toArray
    for (i <- 0 to allVars.length - 1)
      for (j <- i + 1 to allVars.length - 1)
        if (allVars(i) == allVars(j))
         throw new Exception(s"Variable ${allVars(i)} is used twice in the transition system")

    // Every state variable must have an entry in the transition map
    stateIds.find(v => !transition.contains(v)) match {
      case Some(v) =>
        throw new Exception(s"State variable $v must have an entry in the transition map")
      case None => ()
    }

    // Every auxiliary variable must have an entry in the transition map
    auxiliaryIds.find(v => !transition.contains(v)) match {
      case Some(v) =>
        throw new Exception(s"Auxiliary variable $v must have an entry in the transition map")
      case None => ()
    }

    // Every auxiliary variable computation can only depend on previous auxiliary
    // variables (and the state and input variables)
    val auxiliaryArray = auxiliaryIds.toArray
    for (i <- 0 to auxiliaryArray.length - 1) {
      transition(auxiliaryArray(i)).fv.find(id =>
        !stateIds.contains(id) &&
        !inputIds.contains(id) &&
        !auxiliaryIds.take(i).contains(id)
      ) match {
        case None => ()
        case Some(id) =>
          throw new Exception(s"The computation of ${auxiliaryArray(i)} is not allowed to depend on $id")
      }
    }
  }

  def next(q: State, a: Input): State = {
    val env = evalVars(q ++ a, auxiliaryIds, transition)
    evalVarsWithAcc(env, Map(), stateIds, transition).collect {
      case (k: StateIdentifier, v) => (k, v)
    }
  }

  // Check whether there exist:
  // * a state `s0` that makes `from` evaluate to `true`,
  // * a state `sk` that makes `invariant` evaluate to `false`,
  // * a sequence of intermediary states `s1,...,s_{k-1}`
  // * a sequence of inputs `a0,...,a_{k-1}` such that for all i, next(s_i, a_i) = s_{i+1}
  //
  // When this is the case, the function returns
  // Some(Seq(s0, ..., sk), Seq(a0, ..., a_{k-1})), and
  // None otherwise.
  def checkErrorReachability(from: Expr, error: Expr, k: Int): Option[(Seq[State], Seq[Input])] = {
    ??? // TODO (Problem 3)
  }

  def checkErrorReachabilityAtMost(from: Expr, invariant: Expr, k: Int): Option[(Seq[State], Seq[Input])] = {
    def loop(i: Int): Option[(Seq[State], Seq[Input])] = {
      if (i <= k) checkErrorReachability(from, invariant, i) orElse loop(i + 1)
      else None
    }
    loop(0)
  }

  def checkReachability(from: State, errorState: State, k: Int): Option[(Seq[State], Seq[Input])] = {
    checkErrorReachability(stateToExpr(from), stateToExpr(errorState), k)
  }

  def checkReachabilityAtMost(from: State, errorState: State, k: Int): Option[(Seq[State], Seq[Input])] = {
    checkErrorReachabilityAtMost(stateToExpr(from), stateToExpr(errorState), k)
  }

  def dump(fileName: String, hlPath: Option[Path] = None, toPNG: Boolean = true) = {

    // hlPath foreach { case (path, finalState) =>
    //   println("Printing with highlight...\n")
    //   println("Path: ")
    //   path foreach { case (state, input) =>
    //     println(s" -\\ ")
    //     println(s"  | State: $state")
    //     println(s"  | Input: $input")
    //   }
    //   println(s"\nFinal state:\n - $finalState")
    //   println()
    // }

    val baseName = fileName.replaceAll("\\.[^.]*$", "")
    val pngName = baseName + ".png"
    val dotName = baseName + ".dot"
    val out = new java.io.PrintStream(dotName)

    val allStates: Set[State] = enumerate(stateIds)
    val allInputs: Set[Input] = enumerate(inputIds)
    val intOfState: Map[State, Int] = allStates.zipWithIndex.toMap

    out.println("digraph dfa {")
    for ((q, i) <- intOfState) {
      if (hlPath.exists(path => path._2 == q || path._1.toSet.exists(_._1 == q))) {
        out.println(s"""\t$i [label = "${mapToString(q)}", fillcolor = green3, style = filled];""")
      } else {
        out.println(s"""\t$i [label = "${mapToString(q)}"];""")
      }
    }
    out.println()

    for ((p, i) <- intOfState) {
      for (input <- allInputs) {
        val q = next(p, input)
        if (hlPath.exists(path => path._1.toSet.contains((p, input)))) {
          out.println(s"""\t$i -> ${intOfState(q)} [ label = "${mapToString(input)}", color = crimson, fontcolor = crimson ];""")
        } else {
          out.println(s"""\t$i -> ${intOfState(q)} [ label = "${mapToString(input)}" ];""")
        }
      }
    }

    out.println("}")
    out.close
    if (toPNG)
      "dot -Tpng -o %s %s".format(pngName,dotName) !!;
  }

}

object TransitionSystem {
  type State = Map[StateIdentifier, Boolean]
  type Input = Map[InputIdentifier, Boolean]
  type Path = (Seq[(State, Input)], State)

  // After evaluating the expression of a given variable, this function saves
  // the result in an accumulator which is not visible when interpreting later
  // variables
  def evalVarsWithAcc(
    env: Map[Identifier, Boolean],
    acc: Map[Identifier, Boolean],
    vars: Seq[Identifier],
    transition: Map[Identifier, Expr]
  ): Map[Identifier, Boolean] = vars match {
    case Seq() => acc
    case x +: xs => evalVarsWithAcc(env, acc.updated(x, transition(x).interpret(env)), xs, transition)
  }

  // After evaluating the expression of a given variable, this function saves
  // the result in the environment so that it can be used for later variables
  def evalVars(
    env: Map[Identifier, Boolean],
    vars: Seq[Identifier],
    transition: Map[Identifier, Expr]
  ): Map[Identifier, Boolean] = vars match {
    case Seq() => env
    case x +: xs => evalVars(env.updated(x, transition(x).interpret(env)), xs, transition)
  }

  // Returns an expression whose only satisfying assignment is `s`
  def stateToExpr(s: State): Expr = {
    and(s.map { case (id, b) => Equal(Var(id), BooleanLiteral(b)): Expr })
  }

  // Returns a random transition system
  def random(stateIds: Seq[StateIdentifier], auxiliaryIds: Seq[AuxiliaryIdentifier], inputIds: Seq[InputIdentifier]): TransitionSystem = {
    val computationTransition: Map[Identifier, Expr] = auxiliaryIds.foldLeft((stateIds ++ inputIds, Map[Identifier, Expr]())){
      case ((fv, acc), id) =>
        (fv :+ id, acc.updated(id, Expr.random(fv.toArray)))
    }._2
    val varArray: Array[Identifier] = (stateIds ++ inputIds ++ auxiliaryIds).toArray
    val stateTransition: Map[Identifier, Expr] = stateIds.map(id => id -> Expr.random(varArray)).toMap
    TransitionSystem(stateIds, auxiliaryIds, inputIds, computationTransition ++ stateTransition)
  }

}
