import TransitionSystem._
import Expr._

object Adder {
  def accId(i: Int): StateIdentifier = StateIdentifier(s"acc$i")
  def counterId(i: Int): StateIdentifier = StateIdentifier(s"counter$i")
  val overflow: StateIdentifier = StateIdentifier("overflow")
  def carryId(i: Int): AuxiliaryIdentifier = AuxiliaryIdentifier(s"carry$i")
  def inputId(i: Int): InputIdentifier = InputIdentifier(s"input$i")

  val invariant: Expr = Not(Var(overflow))

  // The adder is a transition system.
  // * It maintains an accumulator with `n` bits
  // * It takes as input `m` bits (with m <= n) which are to be added to the accumulator
  // * It maintains a counter with `l` bits, counting the number of inputs that have been received
  // * When the counter is full (all booleans are true), the adder just loops in the same state and stops accumulating inputs
  // * It maintains a boolean variable `overflow` which is set to true when the accumulator overflows
  def adder(n: Int, m: Int, l: Int): TransitionSystem = {
    require(m <= n)
    val accIds: Seq[StateIdentifier] = (0 to n-1).map(accId)
    val counterIds: Seq[StateIdentifier] = (0 to l-1).map(counterId)
    val carryIds: Seq[AuxiliaryIdentifier] = (0 to n).map(carryId)
    val inputIds: Seq[InputIdentifier] = (0 to m-1).map(inputId)
    val stateIds: Seq[StateIdentifier] = accIds ++ counterIds :+ overflow

    val computeCarry: Map[Identifier, Expr] =
      (1 to n).map { i =>
        carryId(i) ->
          ??? // TODO (Problem 1)
      }.toMap[Identifier, Expr].updated(
        carryId(0),
        BooleanLiteral(false)
      )

    val computeAcc: Map[Identifier, Expr] =
      (0 to n-1).map { i =>
        accId(i) ->
          ??? // TODO (Problem 1)
      }.toMap

    val computeCount: Map[Identifier, Expr] =
      (0 to l-1).map { i =>
        counterId(i) ->
          ??? // TODO (Problem 1)
      }.toMap

    val computeOverflow: Map[Identifier, Expr] = Map(
      overflow -> ??? // TODO (Problem 1)
    )

    TransitionSystem(
      stateIds,
      carryIds,
      inputIds,
      computeAcc ++ computeCarry ++ computeCount ++ computeOverflow
    )
  }

  def mapToBigInt[T](n: Int, m: Map[T, Boolean], f: Int => T): BigInt = {
    BigInt((n-1 to 0 by -1).map(i => if (m(f(i))) '1' else '0').mkString, 2)
  }

  def overflowToBigInt(state: State): BigInt = mapToBigInt(1, state, _ => overflow)
  def stateToBigInt(n: Int, state: State): BigInt = mapToBigInt(n, state, accId)
  def inputToBigInt(n: Int, input: Input): BigInt = mapToBigInt(n, input, inputId)

  def intToMap[T](n: Int, a: Int, f: Int => T): Map[T, Boolean] = {
    BigInt(a).toString(2).reverse.padTo(n, '0').zipWithIndex.map { case (c,i) =>
      (f(i), c == '1')
    }.toMap
  }

  def interactiveTest(n: Int, m: Int, l: Int) = {
    val adder: TransitionSystem = Adder.adder(n, m, l)
    adder.checkWF

    def loop(state: State): Unit = {
      println(stateToString(n, l, state))
      println()
      val i = scala.io.StdIn.readLine("Enter 'STOP' or an input: ").trim
      println()
      if (i != "STOP") {
        val input = intToMap(m, i.toInt, inputId)
        loop(adder.next(state, input))
      }
    }
    loop(adder.stateIds.map(id => id -> false).toMap)
  }

  def stateToString(n: Int, l: Int, s: State): String = {
    val overflowBit = overflowToBigInt(s.filter(_._1 == overflow))
    val accumulator = stateToBigInt(n, s.filter(_._1.id.startsWith("acc")))
    val counter = mapToBigInt(l, s.filter(_._1.id.startsWith("counter")), counterId)

    s"Accumulator: $accumulator (${accumulator.toString(2).reverse.padTo(n, '0')})\n" +
    s"Counter:     $counter (${counter.toString(2).reverse.padTo(n, '0')})\n" +
    s"Overflow:    $overflowBit"
  }

  def checkNoOverflow(n: Int, m: Int, l: Int): Unit = {
    val adder = Adder.adder(n, m, l)
    val initial: Expr = stateToExpr(adder.stateIds.map(id => id -> false).toMap)
    val maxSteps: Int = BigInt(2).pow(m).toInt
    adder.checkErrorReachabilityAtMost(initial, Not(Adder.invariant), maxSteps) match {
      case None =>
        println(" => No overflow")

      case Some((sSeq, aSeq)) =>
        for ((s, a) <- sSeq.zip(aSeq)) {
          println("State:")
          println(stateToString(n, l, s))
          println()
          println("Input: " + inputToBigInt(m, a))
          println()
        }
        println("Final state:")
        println(stateToString(n, l, sSeq.last))
        println()
        println(" => Overflow")
    }
  }

  def inductiveInvariant(n: Int, m: Int, l: Int): Expr = {
    ??? // TODO (Problem 3)
  }

  def checkInductiveInvariant(n: Int, m: Int, l: Int): Unit = {
    checkInductiveInvariant(inductiveInvariant(n, m, l), n, m, l)
  }

  def checkInductiveInvariant(inv: Expr, n: Int, m: Int, l: Int): Unit = {
    val adder = Adder.adder(n, m, l)
    val initial: Expr = stateToExpr(adder.stateIds.map(id => id -> false).toMap)

    // println(s"Checking that $inv is inductive")

    adder.checkErrorReachability(initial, Not(inv), 0) match {
      case None =>
        println(" => OK: The invariant is initially valid")
      case Some((sSeq, aSeq)) =>
        println(" => ERROR: The invariant is not initially valid")
    }

    adder.checkErrorReachability(inv, Not(Adder.invariant), 0) match {
      case None =>
        println(" => OK: Your inductive invariant implies the (no overflow) invariant")
      case Some((sSeq, aSeq)) =>
        println(" => ERROR:  Your inductive invariant does not imply the (no overflow) invariant")
    }

    adder.checkErrorReachability(inv, Not(inv), 1) match {
      case None =>
        println(" => OK: The invariant is inductive")

      case Some((sSeq, aSeq)) =>
        for ((s, a) <- sSeq.zip(aSeq)) {
          println("State:")
          println(stateToString(n, l, s))
          println()
          println("Input: " + inputToBigInt(m, a))
          println()
        }
        println("Final state:")
        println(stateToString(n, l, sSeq.last))
        println()
        println(" => ERROR: The invariant is not inductive")
    }
  }
}
