// Countdown Game using the Powerset Function
//============================================


def powerset(xs: Set[Int]) : Set[Set[Int]] = {
  if (xs == Set()) Set(Set())
  else {
    val ps = powerset(xs.tail)  
    ps ++ ps.map(_ + xs.head)
  }
}  

powerset(Set(1,2,3)).mkString("\n")

// proper subsets
def psubsets(xs: Set[Int]) = 
  powerset(xs) -- Set(Set(), xs) 

psubsets(Set(1,2,3)).mkString("\n")

// pairs of subsets and their "complement"
def splits(xs: Set[Int]) : Set[(Set[Int], Set[Int])] =
  psubsets(xs).map(s => (s, xs -- s))

splits(Set(1,2,3,4)).mkString("\n")


// ususal trees representing infix notation for expressions
enum Tree {
  case Num(i: Int)
  case Add(l: Tree, r: Tree)
  case Mul(l: Tree, r: Tree)
  case Sub(l: Tree, r: Tree)
  case Div(l: Tree, r: Tree)
}
import Tree._

//pretty printing for trees
def pp(tr: Tree) : String = tr match {
  case Num(n) => s"$n"
  case Add(l, r) => s"(${pp(l)} + ${pp(r)})"
  case Mul(l, r) => s"(${pp(l)} * ${pp(r)})"
  case Sub(l, r) => s"(${pp(l)} - ${pp(r)})"
  case Div(l, r) => s"(${pp(l)} / ${pp(r)})"
}

// evaluating a tree - might fail when dividing by 0 
// (the for-notation makes it easy to deal with Options)
def eval(tr: Tree) : Option[Int] = tr match {
  case Num(n) => Some(n)
  case Add(l, r) => 
    for (ln <- eval(l); rn <- eval(r)) yield ln + rn
  case Mul(l, r) => 
    for (ln <- eval(l); rn <- eval(r)) yield ln * rn 
  case Sub(l, r) => 
    for (ln <- eval(l); rn <- eval(r); if rn <= ln) yield ln - rn
  case Div(l, r) => 
    for (ln <- eval(l); rn <- eval(r); if rn != 0) 
      yield ln / rn 
}


// simple-minded generation of nearly all possible expressions
// (the nums argument determines the set of numbers over which
//  the expressions are generated)
def gen1(nums: Set[Int]) : Set[Tree] = nums.size match {
  case 0 => Set()
  case 1 => Set(Num(nums.head))
  case 2 => {
    val ln = Num(nums.head)
    val rn = Num(nums.tail.head)
    Set(Add(ln, rn), Mul(ln, rn),
        Sub(ln, rn), Sub(rn, ln),
        Div(ln, rn), Div(rn, ln))
  }
  case n => {
    val res = 
      for ((ls, rs) <- splits(nums);
            lt <- gen1(ls);
            rt <- gen1(rs)) yield {
            Set(Add(lt, rt), Mul(lt, rt),
                Sub(lt, rt), Sub(rt, lt),
                Div(lt, rt), Div(rt, lt))
           }
    res.flatten
  }
}


// some testcases
gen1(Set(1))
gen1(Set(1, 2)).mkString("\n")
gen1(Set(1, 2, 3)).map(pp).mkString("\n")
gen1(Set(1, 2, 3)).map(tr => s"${pp(tr)} = ${eval(tr)}").mkString("\n")

gen1(Set(1, 2, 3)).size             // => 168
gen1(Set(1, 2, 3, 4, 5, 6)).size    // => 26951040

/*
    It is clear that gen1 generates too many expressions
    to be fast overall.

    An easy fix is to not generate improper Subs and 
    Divs when they are at the leaves.
*/


def gen2(nums: Set[Int]) : Set[Tree] =  nums.size match {
  case 0 => Set()
  case 1 => Set(Num(nums.head))
  case 2 => {
    val l = nums.head
    val r = nums.tail.head
    Set(Add(Num(l), Num(r)), 
        Mul(Num(l), Num(r)))
        ++ Option.when(l <= r)(Sub(Num(r), Num(l)))
        ++ Option.when(l > r)(Sub(Num(l), Num(r)))
        ++ Option.when(r > 0 && l % r == 0)(Div(Num(l), Num(r)))
        ++ Option.when(l > 0 && r % l == 0)(Div(Num(r), Num(l)))
  }
  case xs => {
    val res = 
      for ((lspls, rspls) <- splits(nums);
           lt <- gen2(lspls); 
           rt <- gen2(rspls)) yield {
        Set(Add(lt, rt), Sub(lt, rt),
            Mul(lt, rt), Div(lt, rt))
    } 
    res.flatten
  }
}

gen2(Set(1, 2, 3)).size             // => 68
gen2(Set(1, 2, 3, 4, 5, 6)).size    // => 6251936

// OK, the numbers decreased in gen2 as it does 
// not generate leaves like (2 - 3) and (4 / 3).
// It might still generate such expressions "higher
// up" though, for example (1 + 2) - (4 + 3)

println(gen2(Set(1,2,3,4)).map(pp).mkString("\n"))

// before taking any time measure, it might be good
// to check that no "essential" expression has been
// lost by the optimisation in gen2...some eyeballing

gen1(Set(1,2,3,4)).map(eval) == gen2(Set(1,2,3,4)).map(eval)
gen1(Set(1,2,3,4,5,6)).map(eval) == gen2(Set(1,2,3,4,5,6)).map(eval)


// lets start some timing

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

// check to reach a target
def check(xs: Set[Int], target: Int) =
  gen2(xs).find(eval(_) == Some(target))

// example 1
check(Set(50, 5, 4, 9, 10, 8), 560).foreach { sol =>
  println(s"${pp(sol)} => ${eval(sol)}")
}
// => ((50 + ((4 / (10 - 9)) * 5)) * 8) => Some(560)

time_needed(1, check(Set(50, 5, 4, 9, 10, 8), 560))
// => ~14 secs


// example 2
check(Set(25, 5, 2, 10, 7, 1), 986).foreach { sol =>
  println(s"${pp(sol)} => ${eval(sol)}")
}

time_needed(1, check(Set(25, 5, 2, 10, 7, 1), 986))
// => ~15 secs

// example 3 (unsolvable)
check(Set(25, 5, 2, 10, 7, 1), -1)
time_needed(1, check(Set(25, 5, 2, 10, 7, 1), -1))
// => ~22 secs

// example 4
check(Set(100, 25, 75, 50, 7, 10), 360).foreach { sol =>
  println(s"${pp(sol)} => ${eval(sol)}")
}
time_needed(1, check(Set(100, 25, 75, 50, 7, 10), 360))
// => ~14 secs

/*
  Twenty-two seconds in the worst case...this does not yet 
  look competitive enough.

  Lets generate the expression together with the result 
  and restrict the number of expressions in this way.
*/


def gen3(nums: Set[Int]) : Set[(Tree, Int)] =  nums.size match {
  case 0 => Set()
  case 1 => Set((Num(nums.head), nums.head))
  case xs => {
    val res =
      for ((lspls, rspls) <- splits(nums);
           (lt, ln) <- gen3(lspls); 
           (rt, rn) <- gen3(rspls)) yield {
        Set((Add(lt, rt), ln + rn),
            (Mul(lt, rt), ln * rn))
        ++ Option.when(ln <= rn)((Sub(rt, lt), rn - ln)) 
        ++ Option.when(ln > rn)((Sub(lt, rt), ln - rn))
        ++ Option.when(rn > 0 && ln % rn == 0)((Div(lt, rt), ln / rn))
        ++ Option.when(ln > 0 && rn % ln == 0)((Div(rt, lt), rn / ln))
    } 
    res.flatten
  }
}

// eyeballing that we did not lose any solutions...
gen2(Set(1,2,3,4)).map(eval).flatten == gen3(Set(1,2,3,4)).map(_._2)
gen2(Set(1,2,3,4,5,6)).map(eval).flatten == gen3(Set(1,2,3,4,5,6)).map(_._2)

// the number of generated expressions
gen3(Set(1, 2, 3)).size             // => 104
gen3(Set(1, 2, 3))
gen3(Set(1, 2, 3, 4, 5, 6)).size    // => 5092300

// ...while this version does not "optimise" the leaves
// as in gen2, the number of generated expressions grows
// slower

def check2(xs: Set[Int], target: Int) =
  gen3(xs).find(_._2 == target)


// example 1
time_needed(1, check2(Set(50, 5, 4, 9, 10, 8), 560))
// => ~14 secs

// example 2
time_needed(1, check2(Set(25, 5, 2, 10, 7, 1), 986))
// => ~18 secs

// example 3 (unsolvable)
time_needed(1, check2(Set(25, 5, 2, 10, 7, 1), -1))
// => ~19 secs

// example 4
time_needed(1, check2(Set(100, 25, 75, 50, 7, 10), 360))
// => ~16 secs

// ...we are getting better, but not yet competetive enough.
// 
// The problem is that splits generates for sets say {1,2,3,4}
// the splits ({1,2}, {3,4}) and ({3,4}, {1,2}). This means that
// we consider terms (1 + 2) * (3 + 4) and (3 + 4) * (1 + 2) as
// separate candidates. We can avoid this duplication by returning
// sets of sets of numbers, like 

// pairs of subsets and their "complement"
def splits2(xs: Set[Int]) : Set[Set[Set[Int]]] =
  psubsets(xs).map(s => Set(s, xs -- s))

splits(Set(1,2,3,4)).mkString("\n")
splits2(Set(1,2,3,4)).mkString("\n")

def gen4(nums: Set[Int]) : Set[(Tree, Int)] =  nums.size match {
  case 0 => Set()
  case 1 => Set((Num(nums.head), nums.head))
  case xs => {
    val res =
      for (spls <- splits2(nums);
           (lt, ln) <- gen4(spls.head); 
           (rt, rn) <- gen4(spls.tail.head)) yield {
        Set((Add(lt, rt), ln + rn),
            (Mul(lt, rt), ln * rn))
        ++ Option.when(ln <= rn)((Sub(rt, lt), rn - ln)) 
        ++ Option.when(ln > rn)((Sub(lt, rt), ln - rn)) 
        ++ Option.when(rn > 0 && ln % rn == 0)((Div(lt, rt), ln / rn))
        ++ Option.when(ln > 0 && rn % ln == 0)((Div(rt, lt), rn / ln)) 
    } 
    res.flatten
  }
}

// eyeballing that we did not lose any solutions...
gen4(Set(1,2,3,4)).map(_._2) == gen3(Set(1,2,3,4)).map(_._2)
gen4(Set(1,2,3,4,5,6)).map(_._2) == gen3(Set(1,2,3,4,5,6)).map(_._2)


gen4(Set(1, 2, 3)).size             // => 43
gen4(Set(1, 2, 3, 4))
gen4(Set(1, 2, 3, 4, 5, 6)).size    // => 550218

def check3(xs: Set[Int], target: Int) =
  gen4(xs).find(_._2 == target)

// example 1
check3(Set(50, 5, 4, 9, 10, 8), 560).foreach { sol =>
  println(s"${pp(sol._1)} => ${eval(sol._1)}")
}
// (((10 - 5) * (9 * 8)) + (4 * 50)) => Some(560)

time_needed(1, check3(Set(50, 5, 4, 9, 10, 8), 560))
// => ~1 sec


// example 2
check3(Set(25, 5, 2, 10, 7, 1), 986).foreach { sol =>
  println(s"${pp(sol._1)} => ${eval(sol._1)}")
}

time_needed(1, check3(Set(25, 5, 2, 10, 7, 1), 986))
// => ~1 sec

// example 3 (unsolvable)
check3(Set(25, 5, 2, 10, 7, 1), -1)
time_needed(1, check3(Set(25, 5, 2, 10, 7, 1), -1))
// => ~2 secs

// example 4
check3(Set(100, 25, 75, 50, 7, 10), 360).foreach { sol =>
  println(s"${pp(sol._1)} => ${eval(sol._1)}")
}
time_needed(1, check3(Set(100, 25, 75, 50, 7, 10), 360))
// => ~1 secs



time_needed(1, check3(Set(50, 5, 4, 9, 10, 8), 560))
time_needed(1, check3(Set(25, 5, 2, 10, 7, 1), 986))
time_needed(1, check3(Set(25, 5, 2, 10, 7, 1), -1))
time_needed(1, check3(Set(100, 25, 75, 50, 7, 10), 360))