// 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 io.Source
import scala.util._

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

// (6)
def jumpRight(prog: String, pc: Int, level: Int) : Int = {
    if (pc >= prog.length) prog.length
    else if (prog(pc) == '[') jumpRight(prog, pc + 1, level + 1)
    else if (prog(pc) == ']' && level == 0) pc + 1
    else if (prog(pc) == ']') jumpRight(prog, pc + 1, level - 1)
    else jumpRight(prog, pc + 1, level)
}

def jtable(pg: String) : Map[Int, Int] = {
    val pairs = for {
        i <- 0 until pg.length
        if pg.charAt(i) == '['
        j = jumpRight(pg, i+1, 0)
    } yield (i, j)
    pairs.flatMap { case (i, j) => 
        List((i, j), (j-1, i+1))
    }.toMap
}

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

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

def compute2(pg: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = {
  if (pc >= pg.length) mem
  else {
      val (npc, nmp, nmem) = pg(pc) match {
          case '>' => (pc + 1, mp + 1, mem)
          case '<' => (pc + 1, mp - 1, mem)
          case '+' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) + 1)))
          case '-' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) - 1)))
          case '.' => {print(mem.getOrElse(mp,0).toChar);(pc + 1, mp, mem)}
          case '[' => if (mem.getOrElse(mp, 0) == 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case ']' => if (mem.getOrElse(mp, 0) != 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case _ => (pc + 1, mp, mem)
      }
      compute2(pg, tb, npc, nmp, nmem)
  }
}

def run2(pg: String, m: Mem = Map()) = 
  compute2(pg, jtable(pg), 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.replaceAll("""[^<>+-.,\[\]]""","").replaceAll("""\[-\]""","0")

def compute3(pg: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = {
    if (pc >= pg.length) mem
    else {
      val (npc, nmp, nmem) = pg(pc) match {
          case '>' => (pc + 1, mp + 1, mem)
          case '<' => (pc + 1, mp - 1, mem)
          case '+' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) + 1)))
          case '-' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) - 1)))
          case '.' => {print(mem.getOrElse(mp,0).toChar);(pc + 1, mp, mem)}
          case '[' => if (mem.getOrElse(mp, 0) == 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case ']' => if (mem.getOrElse(mp, 0) != 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case _ => (pc + 1, mp, mem)
      }
      compute3(pg, tb, npc, nmp, nmem)
    }
}

def run3(pg: String, m: Mem = Map()) = {
  val opt_pg = optimise(pg)
  val jt = jtable(opt_pg)
  compute3(opt_pg, jt, 0, 0, m)
}


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


// (8)  
def combine(s: String): String = ???

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

def compute4(pg: String, tb: Map[Int, Int], pc: Int, mp: Int, mem: Mem) : Mem = {
    if (pc >= pg.length) mem
    else {
      val (npc, nmp, nmem) = pg(pc) match {
          case '>' => (pc + 1, mp + 1, mem)
          case '<' => (pc + 1, mp - 1, mem)
          case '+' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) + 1)))
          case '-' => (pc + 1, mp, mem + (mp -> (mem.getOrElse(mp, 0) - 1)))
          case '.' => {print(mem.getOrElse(mp,0).toChar);(pc + 1, mp, mem)}
          case '[' => if (mem.getOrElse(mp, 0) == 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case ']' => if (mem.getOrElse(mp, 0) != 0) (tb(pc), mp, mem) else (pc + 1, mp, mem)
          case _ => (pc + 1, mp, mem)
      }
      compute3(pg, tb, npc, nmp, nmem)
    }
}

// should call first optimise and then combine on the input string
//
def run4(pg: String, m: Mem = Map()) = {
  val co_opt_pg = combine(optimise(pg))
  val jt = jtable(co_opt_pg)
  compute3(co_opt_pg, jt, 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")))


}





// This template code is subject to copyright 
// by King's College London, 2022. Do not 
// make the template code public in any shape 
// or form, and do not exchange it with other 
// students under any circumstance.