import TransitionSystem._
import Expr._

object Arbiter {
  val idle   = StateIdentifier("Idle")
  val grantA = StateIdentifier("Grant_A")
  val grantB = StateIdentifier("Grant_B")
  val grantC = StateIdentifier("Grant_C")

  val Idle   = Var(idle)
  val GrantA = Var(grantA)
  val GrantB = Var(grantB)
  val GrantC = Var(grantC)

  val reqA = InputIdentifier("Req_A")
  val reqB = InputIdentifier("Req_B")
  val reqC = InputIdentifier("Req_C")

  val ReqA = Var(reqA)
  val ReqB = Var(reqB)
  val ReqC = Var(reqC)

  val stateVars = Seq(idle, grantA, grantB, grantC)
  val inputVars = Seq(reqA, reqB, reqC)

  val initialState: State = Map(
    idle   -> true,
    grantA -> false,
    grantB -> false,
    grantC -> false,
  )

  // TODO: Fill in the invariant (Problem 1.1)
  val invariant: Expr = ???

  // TODO: Write down the input that leads the system to an error state (Problem 1.2)
  // Input: ???

  // TODO: Fix the transition system (Problem 1.3)
  val transition: Map[Identifier, Expr] = Map(
    grantA -> or(
      and(Idle,   ReqA),
      and(GrantA, ReqA),
      and(GrantB, ReqA, ReqB.bar),
      and(GrantC, ReqA, ReqC.bar),
    ),

    grantB -> or(
      and(Idle,   ReqB),
      and(GrantA, ReqA.bar, ReqB),
      and(GrantB, ReqB),
      and(GrantC, ReqA.bar, ReqB, ReqC.bar),
    ),

    grantC -> or(
      and(Idle,   ReqC),
      and(GrantA, ReqA.bar, ReqB.bar, ReqC),
      and(GrantB, ReqA.bar, ReqB.bar, ReqC),
      and(GrantC, ReqC),
    ),

    idle -> or(
      and(Idle,   ReqA.bar, ReqB.bar, ReqC.bar),
      and(GrantA, ReqA.bar, ReqB.bar, ReqC.bar),
      and(GrantB, ReqA.bar, ReqB.bar, ReqC.bar),
      and(GrantC, ReqA.bar, ReqB.bar, ReqC.bar),
    ),
  )

  val arbiter = TransitionSystem(
    stateVars,
    Seq.empty,
    inputVars,
    transition,
  )

  def verify(): Unit = {
    arbiter.dump("arbiter", toPNG = true)

    val initialExpr = stateToExpr(initialState)

    val path = arbiter.checkErrorReachabilityAtMost(initialExpr, Not(invariant), 20)

    path match {
      case None =>
        println(" => Error state not reachable.")

      case Some((sSeq, aSeq)) =>
        for ((state, input) <- sSeq.zip(aSeq)) {
          println(s"State: $state")
          println(s"Input: $input")
          println(s"")
        }
        println(s"Final state: ${sSeq.last}\n")
        println(" => Error state reachable! (see path in arbiter-path.png)")

        arbiter.dump("arbiter-path", Some((sSeq.zip(aSeq), sSeq.last)), toPNG = true)
    }
  }

  def simulate(): Unit = {
    val inputs = Seq(
      Map(reqA -> true,  reqB -> true,  reqC -> true),
      Map(reqA -> true,  reqB -> true,  reqC -> true),
      Map(reqA -> true,  reqB -> false, reqC -> true),
      Map(reqA -> true,  reqB -> false, reqC -> true),
      Map(reqA -> false, reqB -> false, reqC -> true),
    )

    println(s"Initial state: $initialState")

    var state = initialState
    for (input <- inputs) {
      state = arbiter.next(state, input)
      println(s" -\\")
      println(s"  | Input: $input")
      println(s"  | State: $state")
    }
  }

}
