require 'ai4r/neural_network/backpropagation'
require 'prime'
# srand 1

def bits_for(n)
  [n[0],n[1],n[2],n[3],n[4],n[5],n[6],n[7]]
end

def train(net, n)
  net.train bits_for(n), [n.prime? ? 1 : 0]
end

loop do
  puts "Training the network, please wait.";
  max = (2**8)-1

  net = Ai4r::NeuralNetwork::Backpropagation.new([8, 9, 1]);

  2001.times do |i|
    errors = 0.upto(max).to_a.shuffle.map { |n| [n, train(net, n)] }
    next unless i % 200 == 0
    average = errors.map(&:last).inject(:+) / errors.length
    puts "Error for run #{i}: #{average.inspect}"
    next if i < 600
    # retrain = errors.select { |n, err| err > 0.2 }
    retrain = 0.upto(max).reject { |n| (0.5 < net.eval(bits_for n).first) == n.prime? }
    puts "  retraining on #{retrain.inspect}"
    500.times { retrain.each { |n| train net, n } }
  end

  puts "Test data"
  num_correct   = 0
  num_attempted = 0
  0.upto max do |n|
    num_attempted += 1
    result   = net.eval(bits_for n).first
    actual   = (0.5 <= result)
    expected = n.prime?
    (expected == actual) ? (num_correct += 1) : printf("%3d \e[31mreal:%-5s ai:%-5s\e[0m (%f)\n", n, expected, actual, result)
  end

  puts
  puts "SUMMARY: #{num_correct}/#{num_attempted} (#{100.0*num_correct/num_attempted}%)"
  if num_attempted == num_correct
    require "pry"
    binding.pry
  end
  puts "-----------------------------"
end