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


local data = {}
-- local csvFile = io.open('HistoricalData_CRNC.csv')
local csvFile = io.open('DJI_full_1min.txt')
local header = true
local mapping = {}
for line in csvFile:lines() do
  if header then
    for lineValue in string.gmatch(line, '[^,]+') do
      table.insert(mapping, lineValue)
    end
    header = false
  else
    local tbl = {}
    local i = 1
    for lineValue in string.gmatch(line, '[^,]+') do
      tbl[mapping[i]] = tonumber(string.match(lineValue, '[^$]+'))
      i = i+1
    end
    table.insert(data, tbl)
  end
end


local metaproduction = {}
metaproduction["<start>"] = {{'<functions>'}}
metaproduction["<functions>"] = { {}, {'<function>', '<functions>'} }
metaproduction["<function>"] = {{ 'defun(', '<arity>', ')\n' }}
metaproduction["<arity>"] = { {'0'}, {'1+', '<arity>'} }

local metagrammar = Grammar.createGrammar(metaproduction, {})

function baseProduction()
  local production = {}
  production["<start>"]={{ 'return ','<expr>'}}
  production["<expr>"]={ 
    {'x'}, {'<number>'}, 
    -- {'<data>'}, {'<number>'}, 
    -- {'<data>'}, {'<number>'}, 
    {'(','<expr>','<op>','<expr>',')'}, {'<func>','(','<expr>',')'} 
  }
  production["<op>"]={ {'+'}, {'-'}, {'*'}, {'/'} }
  production["<func>"]={ {'math.log'}, {'math.exp'}, {'math.sin'}, {'math.cos'}, {'math.tan'}, }
  production["<digit>"]={ {'0'}, {'1'}, {'2'}, {'3'}, {'4'}, {'5'}, {'6'}, {'7'}, {'8'}, {'9'} }
  production["<lowdigit>"]={ {'0'}, {'1'}, {'2'} }
  production["<data>"]={ 
    {'high(','<date>',')'}, {'low(','<date>',')'}, 
    -- {'volume(','<date>',')'}, 
    {'open(','<date>',')'}, {'close(','<date>',')'}
  }
  production["<date>"]={ {'<lowdigit>','<digit>', '<digit>'} }
  production["<number>"]={ {'0.','<digit>','<digit>','<digit>'} }
  return production
end


function deepcopy(tbl)
  local copy = {}
  for k, v in pairs(tbl) do
    if type(v) == 'table' then
      copy[k] = deepcopy(v)
    else
      copy[k] = v
    end
  end
  return copy
end


function expandGenome(genome, mainArgs, env)
  mainArgs = mainArgs or ''
  env = env or _G
  local meta, offset = metagrammar:expand(genome)
  local metafn = assert(loadstring('return function(defun) \n' .. meta .. 'end'))()
  local fnTbl = {}
  local defun = function(arity) table.insert(fnTbl, arity) end
  metafn(defun)
  local production = baseProduction()
  for fn, arity in ipairs(fnTbl) do
    local fnName = string.format("fn%d", fn)
    local args = {}
    for i=1,arity do
      local argName = string.format('arg%d', i)
      table.insert(args, argName)
    end

    -- create function body
    local fnProd = deepcopy(production)
    local fnPrefix = '  local ' .. fnName .. ' = function('
    for i=1,#args do
      table.insert(fnProd["<expr>"], {args[i]})
      fnPrefix = fnPrefix .. args[i]
      if i ~= #args then 
        fnPrefix = fnPrefix .. ', ' 
      end
    end
    fnPrefix = fnPrefix .. ') '
    local fnGrammar = Grammar.createGrammar(fnProd, {})
    local fnStr
    fnStr, offset = fnGrammar:expand(genome, offset)
    fnStr = fnPrefix .. fnStr .. ' end'
    fnTbl[fn] = fnStr

    -- add function to production rules
    local fnExpansion = { fnName, '(' }
    for i=1,arity do
      table.insert(fnExpansion, '<expr>')
      if i ~= arity then
        table.insert(fnExpansion, ', ')
      end
    end
    table.insert(fnExpansion, ')')
    table.insert(production["<expr>"], fnExpansion)
  end

  -- create main production and compile
  local grammar = Grammar.createGrammar(production, {})
  local mainStr = grammar:expand(genome, offset)
  local funcs = ''
  for _, fnStr in ipairs(fnTbl) do
    funcs = funcs .. fnStr .. '\n'
  end
  mainStr = 'return function(' .. mainArgs .. ')\n' .. funcs .. '  ' .. mainStr .. '\nend'
  return mainStr
end


function compileGenome(genome, mainArgs, env)
  local mainStr = expandGenome(genome, mainArgs, env)
  -- print(mainStr)
  local f = assert(loadstring(mainStr))()
  setfenv(f, env or {})
  return f

end


math.randomseed(0)
-- math.randomseed(1735920315)
--local t = os.time()
--print(t)
--math.randomseed(t)

function randomGenome()
  local len = math.ceil(math.random() * 3000)
  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 target = function(x) return (x^4) + (x^3) + (x^2) + x end
  local err, f = pcall(compileGenome, genome, 'x')
  if err == false then
    return -1e6
  end
  local cost = 0
  for n=1,1000 do
    local env = {}
    local future = 24*60
    local range = (#data - future - 300)
    local start = math.ceil((range * math.random())) + future
    local dataView = function(offset) 
      -- 1 <= k <= #data
      -- 1 <= 1 + start - future <= #data
      -- 1 <= 1 + start + 299 <= #delta
      return data[1 + start + offset]
    end
    env.close = function(offset) return dataView(offset).close end
    env.volume = function(offset) return dataView(offset).volume end
    env.open = function(offset) return dataView(offset).open end
    env.high = function(offset) return dataView(offset).high end
    env.low = function(offset) return dataView(offset).low end
    env.math = math
    -- print(#data, start)
    -- print(data[2].volume, env.volume(-future), env.volume(299))
    setfenv(f, env)
    local x = (2*math.random()) - 1
    ok, y = pcall(f, x)
    -- if ok and type(y) == 'table' then for k,v in pairs(y) do print(k, v) end end
    if (not ok) or (y == nil) or (y ~= y) then
      -- print('fail', y)
      -- print(expandGenome(genome))
      cost = cost + 100
    else
      -- local delta = env.high(-future) - env.high(0)
      local delta = target(x)
      cost = cost + (0.01 * (y - delta)^2)
    end
  end
  return -(cost)
end

local ga = GA.createGA(randomPop(), {
  evaluate=evaluate,
  crossover=function(a, b)
    local pointA = math.floor(#a * math.random())
    local pointB = math.ceil(#b * math.random())
    local new = {}
    for i=1,pointA do
      table.insert(new, a[i])
    end
    for i=pointB,#b do
      table.insert(new, b[i])
    end
    return new
  end,
  crossoverWeight = 10,
  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,
  mutationWeight = 10,
  reproductionWeight = 0.1,
})


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

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