local Grammar = require 'grammar'
local GA = require 'ga'


local production = {}
production["<start>"]={{'<expr>'}}
production["<expr>"]={ {'<var>'}, {'(', '<expr>', '<op>', '<expr>', ')'} }
production["<op>"]={ {'+'}, {'-'}, {'*'}, {'/'} }
production["<var>"]={ {'x'} }


local grammar = Grammar.createGrammar(production)


math.randomseed(0)


function randomGenome()
  local len = math.ceil(math.random() * 256)
  local genome = {}
  for i=1,len do
    table.insert(genome, math.floor(256 * math.random()))
  end
  return genome
end


function randomPop()
  local pop = {}
  for n=1,1000 do
    table.insert(pop, randomGenome())
  end
  return pop
end

function evaluate(genome)
  local s = 'return function(x) return ' .. grammar:expand(genome) .. ' end'
  local f = assert(loadstring(s))()
  local target = function(x) return (x*x*x*x) + (x*x*x) + (x*x) + (x) end
  local cost = 0
  for n=0,100 do
    x = (n-50)/50
    ok, y = pcall(f, x)
    if (not ok) or (y ~= y) then
      cost = cost + 100
    else
      cost = cost + math.abs(y - target(x))
    end
  end
  return -((0.0001*#s) + cost)
end

local ga = GA.createGA(randomPop(), {
  evaluate=evaluate,
  crossover=function(a, b)
    local crossoverPoint = math.ceil(#a * math.random())
    local new = {}
    for i=1,#b do
      table.insert(new, (i>=crossoverPoint and b[i]) or a[i])
    end
    return new
  end,
  mutate=function(a)
    local new = {}
    for i,v in ipairs(a) do new[i] = v end
    local mutationPoint = math.ceil(#a * math.random())
    new[mutationPoint] = math.floor(256 * math.random())
    return new
  end,
})


for n=1,100 do
  print(ga:step())
  print(grammar:expand(ga.population[ga.bestMember]))
end

-- for k,v in ipairs(ga.population) do
--   print(k, '(', v[1], v[2], v[3], ')', evaluate(v))
-- end