// Main Part 5 about a "Compiler" for the Brainf*** language
//============================================================


object M5b {

// !!! Copy any function you need from file bf.scala !!!
//
// If you need any auxiliary function, feel free to 
// implement it, but do not make any changes to the
// templates below.


// DEBUGGING INFORMATION FOR COMPILERS!!!
//
// Compiler, even real ones, are fiendishly difficult to get
// to produce correct code. One way to debug them is to run
// example programs ``unoptimised''; and then optimised. Does
// the optimised version still produce the same result?


// for timing purposes
def time_needed[T](n: Int, code: => T) = {
  val start = System.nanoTime()
  for (i <- 0 until n) code
  val end = System.nanoTime()
  (end - start)/(n * 1.0e9)
}


type Mem = Map[Int, Int]

import scala.io.Source
import scala.util._

// ADD YOUR CODE BELOW
//======================

def sread(mem: Mem, mp: Int) : Int = mem.getOrElse(mp, 0)

def write(mem: Mem, mp: Int, v: Int) : Mem = mem + (mp -> v)

// (6)
def empty_stack(len : Int, st : List[Int]) : Map[Int, Int] = st match {
  case Nil => Map()
  case n :: tail => empty_stack(len, tail) + (n -> len)
}
def jtable_helper(pg : List[Char], st : List[Int] = List(), index : Int = 0) : Map[Int, Int] = pg match {
  case Nil => empty_stack(pg.length, st)
  case '[' :: tail => jtable_helper(tail, index :: st, index + 1)
  case ']' :: tail => st match {
    case Nil => jtable_helper(tail, st, index + 1) + (index -> -1)
    case n :: stail => jtable_helper(tail, stail, index + 1) + (n -> (index + 1)) + (index -> (n + 1))
  }
  case _ :: tail => jtable_helper(tail, st, index + 1)
}
def jtable(pg: String) : Map[Int, Int] = jtable_helper(pg.toList)

// testcase
//
// jtable("""+++++[->++++++++++<]>--<+++[->>++++++++++<<]>>++<<----------[+>.>.<+<]""")
// =>  Map(69 -> 61, 5 -> 20, 60 -> 70, 27 -> 44, 43 -> 28, 19 -> 6)


def compute2(prog: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = 
    if (pc < 0 || pc >= prog.length) mem else prog.charAt(pc) match {
        case '>' => compute2(prog, tb, pc+1, mp+1, mem)
        case '<' => compute2(prog, tb, pc+1, mp-1, mem)
        case '+' => compute2(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) + 1)))
        case '-' => compute2(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) - 1)))
        case '.' => print(sread(mem, mp).toChar); compute2(prog, tb, pc+1, mp, mem)
        case '[' => if (sread(mem, mp) == 0) compute2(prog, tb, tb.getOrElse(pc, -2), mp, mem)
                    else compute2(prog, tb, pc+1, mp, mem)
        case ']' => if (sread(mem, mp) == 0) compute2(prog, tb, pc+1, mp, mem)
                    else compute2(prog, tb, tb.getOrElse(pc, -2), mp, mem)
        case _ => compute2(prog, tb, pc+1, mp, mem)
}

def run2(prog: String, m: Mem = Map()) = compute2(prog, jtable(prog), 0, 0, m)

// testcases
// time_needed(1, run2(load_bff("benchmark.bf")))
// time_needed(1, run2(load_bff("sierpinski.bf")))



// (7) 

def optimise(s: String) : String =
  s.filter(List('>', '<', '+', '-', '[', ']', '.').contains(_)).replaceAllLiterally("[-]", "0")

def compute3(prog: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = 
  if (pc < 0 || pc >= prog.length) mem else prog.charAt(pc) match {
        case '>' => compute3(prog, tb, pc+1, mp+1, mem)
        case '<' => compute3(prog, tb, pc+1, mp-1, mem)
        case '+' => compute3(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) + 1)))
        case '-' => compute3(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) - 1)))
        case '.' => print(sread(mem, mp).toChar); compute3(prog, tb, pc+1, mp, mem)
        case '[' => if (sread(mem, mp) == 0) compute3(prog, tb, tb.getOrElse(pc, -2), mp, mem)
                    else compute3(prog, tb, pc+1, mp, mem)
        case ']' => if (sread(mem, mp) == 0) compute3(prog, tb, pc+1, mp, mem)
                    else compute3(prog, tb, tb.getOrElse(pc, -2), mp, mem)
        case '0' => compute3(prog, tb, pc+1, mp, mem + (mp -> 0))
        case _ => compute3(prog, tb, pc+1, mp, mem)
}

def run3(prog: String, m: Mem = Map()) = 
  val optimized = optimise(prog)
  compute3(optimized, jtable(optimized), 0, 0, m)


// testcases
//
// optimise(load_bff("benchmark.bf"))          // should have inserted 0's
// optimise(load_bff("mandelbrot.bf")).length  // => 11205 // this is wrong, it's 11203!
// 
// time_needed(1, run3(load_bff("benchmark.bf")))

// (8)
def combine_helper(c : String, s : String) : String =
  val temp1 = s.replaceAllLiterally(c*26, c ++ "Z").replaceAllLiterally(c*25, c ++ "Y").replaceAllLiterally(c*24, c ++ "X").replaceAllLiterally(c*23, c ++ "W")
  val temp2 = temp1.replaceAllLiterally(c*22, c ++ "V").replaceAllLiterally(c*21, c ++ "U").replaceAllLiterally(c*20, c ++ "T").replaceAllLiterally(c*19, c ++ "S")
  val temp3 = temp2.replaceAllLiterally(c*18, c ++ "R").replaceAllLiterally(c*17, c ++ "Q").replaceAllLiterally(c*16, c ++ "P").replaceAllLiterally(c*15, c ++ "O")
  val temp4 = temp3.replaceAllLiterally(c*14, c ++ "N").replaceAllLiterally(c*13, c ++ "M").replaceAllLiterally(c*12, c ++ "L").replaceAllLiterally(c*11, c ++ "K")
  val temp5 = temp4.replaceAllLiterally(c*10, c ++ "J").replaceAllLiterally(c*9, c ++ "I").replaceAllLiterally(c*8, c ++ "H").replaceAllLiterally(c*7, c ++ "G")
  val temp6 = temp5.replaceAllLiterally(c*6, c ++ "F").replaceAllLiterally(c*5, c ++ "E").replaceAllLiterally(c*4, c ++ "D").replaceAllLiterally(c*3, c ++ "C")
  val temp7 = temp6.replaceAllLiterally(c*2, c ++ "B").replaceAllLiterally(c ++ ">", c ++ "A>").replaceAllLiterally(c ++ "<", c ++ "A<")
  val temp8 = temp7.replaceAllLiterally(c ++ "+", c ++ "A+").replaceAllLiterally(c ++ "-", c ++ "A-").replaceAllLiterally(c ++ "[", c ++ "A[")
  temp8.replaceAllLiterally(c ++ "]", c ++ "A]").replaceAllLiterally(c ++ ".", c ++ "A.").replaceAllLiterally(c ++ "0", c ++ "A0")
def combine(s: String) : String = 
  val temp1 = combine_helper(">", s)
  val temp2 = combine_helper("<", temp1)
  val temp3 = combine_helper("+", temp2)
  val temp4 = combine_helper("-", temp3)
  temp4

// testcase
// combine(load_bff("benchmark.bf"))

def compute4(prog: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = if (pc < 0 || pc >= prog.length) mem else prog.charAt(pc) match {
  case '>' => val n = (prog.charAt(pc+1).toInt - 64)
              compute4(prog, tb, pc+1, mp+n, mem)
  case '<' => val n = (prog.charAt(pc+1).toInt - 64)
              compute4(prog, tb, pc+1, mp-n, mem)
  case '+' => val n = (prog.charAt(pc+1).toInt - 64)
              compute4(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) + n)))
  case '-' => val n = (prog.charAt(pc+1).toInt - 64)
              compute4(prog, tb, pc+1, mp, mem + (mp -> (sread(mem, mp) - n)))
  case '.' => print(sread(mem, mp).toChar); compute4(prog, tb, pc+1, mp, mem)
  case '[' => if (sread(mem, mp) == 0) compute4(prog, tb, tb.getOrElse(pc, -2), mp, mem)
              else compute4(prog, tb, pc+1, mp, mem)
  case ']' => if (sread(mem, mp) == 0) compute4(prog, tb, pc+1, mp, mem)
              else compute4(prog, tb, tb.getOrElse(pc, -2), mp, mem)
  case '0' => compute4(prog, tb, pc+1, mp, mem + (mp -> 0))
  case _ => compute4(prog, tb, pc+1, mp, mem)
}

// should call first optimise and then combine on the input string
//
def run4(prog: String, m: Mem = Map()) = 
  val optimized = combine(optimise(prog))
  compute4(optimized, jtable(optimized), 0, 0, m)



// testcases
// combine(optimise(load_bff("benchmark.bf"))) // => """>A+B[<A+M>A-A]<A[[....."""

// testcases (they should now run much faster)
// time_needed(1, run4(load_bff("benchmark.bf")))
// time_needed(1, run4(load_bff("sierpinski.bf"))) 
// time_needed(1, run4(load_bff("mandelbrot.bf")))


}