// 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.