MathMachine

#!/usr/bin/env ruby

require_relative "../lib/z3"

class MathMachine
  def initialize
    @solver = Z3::Solver.new
    @cells     = {}
    @a_source  = {}
    @a_value   = {}
    @b_source  = {}
    @b_value   = {}
    @operation = {}

    # First cell is special
    (1..10).each do |i|
      @cells[i] = cell_var(i)
    end
    @solver.assert @cells[1] == 0

    (2..10).each do |i|
      @a_value[i], @a_source[i] = operand("a", i)
      @b_value[i], @b_source[i] = operand("b", i)
      if i == 2
        @operation[i] = Z3::Int("op#{i}")
        @solver.assert @operation[i] == 4
        @solver.assert Z3::Implies(@operation[i] == 4, 1 == @cells[i])
      elsif i >= 6
        @operation[i] = operation(i, @cells[i-5])
      else
        @operation[i] = operation(i, 0)
      end
    end
  end

  # Just to avoid dumb numbers
  def cell_var(i)
    v = Z3::Int("c#{i}")
    @solver.assert v <= 0x1_0000_0000
    @solver.assert v >= -0x1_0000_0000
    v
  end

  def operand(l,i)
    value  = Z3::Int("#{l}#{i}_value")
    source = Z3::Int("#{l}#{i}_source")
    @solver.assert source >= 1
    @solver.assert source < i
    (1...i).each do |j|
      @solver.assert Z3::Implies(source == j, value == @cells[j])
    end
    [value, source]
  end

  def operation(i,x)
    operation = Z3::Int("op#{i}")
    @solver.assert operation >= 0
    @solver.assert operation <= 3
    @solver.assert Z3::Implies(operation == 0, (@a_value[i] + @b_value[i]) + x == @cells[i])
    @solver.assert Z3::Implies(operation == 1, (@a_value[i] - @b_value[i]) + x == @cells[i])
    @solver.assert Z3::Implies(operation == 2, (@a_value[i] * @b_value[i]) + x == @cells[i])
    @solver.assert Z3::Implies(operation == 3, (@a_value[i] / @b_value[i]) + x == @cells[i])
    @solver.assert Z3::Implies(operation == 3, @b_value[i] != 0)
    operation
  end

  def print_solution!
    puts "Solution:"
    c = @solver.check
    if c == :sat
      m = @solver.model
      (1..10).each do |i|
        print "Cell #{i}: "
        if i != 1
          print "[#{m[@a_source[i]]}] #{m[@a_value[i]]} "
          print "#{ %W[+ - * / **][m[@operation[i]].to_s.to_i] } "
          print "[#{m[@b_source[i]]}] #{m[@b_value[i]]} = "
        end
        puts "#{m[@cells[i]].to_s}"
      end
    else
      puts "* Can't solve the problem (#{c})"
    end
  end

  def solve!
    @solver.assert @cells[10] == 1337
    print_solution!
  end
end

MathMachine.new.solve!

# https://www.sabrefilms.co.uk/revolutionelite/math-machine-basic.php