adaptive_ai = not adaptive_ai and {} or adaptive_ai

dofile("../helper.lua")
dofile("./noise.lua")

trainer = {}

math.randomseed(os.time())

local random = math.random
local helper = adaptive_ai.helper
local calc = adaptive_ai.calc

local ceil = math.ceil
local avg = calc.average
local round = calc.round

local adjacent = helper.adjacent
local conditions = {}
local gen_stats, popul, top_size                                -- condition values
local fill_creature, reset_creature, manage_creature, fill_cell -- condition functions

require "socket"

local function sleep(sec)
    socket.select(nil, nil, sec)
end

local function console_log(s, delay)
    if conditions.output and s then
        print(s)
        if delay then sleep(delay) end
    end
end

-- TERRAIN CONTROL
local function valid(pos)
    if pos.x and pos.x < 1 or pos.x > conditions.dims.x then return false
    elseif pos.z and pos.z < 1 or pos.z > conditions.dims.z then return false
    else return true
    end
end

-- Fill the training area
-- Cell[1] stores the height. Cell[2] is the space to store custom metadata
local function fill_terrain(map)
    console_log("Filling terrain...")

    local p
    local max_x, max_y, max_z = conditions.dims.x, conditions.dims.y, conditions.dims.z
    local x_dimming, z_dimming = max_x/4, max_z/4

    for i=1,max_x do
        for j=1,max_z do
            if map[i][j][1] == -1 then
                p = {x=i, z=j}
                p.y = round((SimplexNoise.Noise2D(p.x/x_dimming, p.z/z_dimming) + 1)*5)
                p.y = p.y + round(SimplexNoise.Noise2D(p.x/2, p.z/2))
                p.y = p.y + round(SimplexNoise.Noise2D(p.x, p.z)/1.5)
                if valid(p) then
                    if p.y > max_y then
                        p.y = max_y
                    elseif p.y < 1 then
                        p.y = 1
                    end
                end
                map[i][j][1] = p.y

                fill_cell(map[i][j])
            end
        end
    end

    console_log("Success filling terrain.\n")
end

local function gen_terrain()
    console_log("Prepare new terrain...")
    local map = {}
    for i=1,40 do
        map[i] = {}
        for j=1,40 do
            map[i][j] = {-1}
        end
    end

    SimplexNoise.seedP(os.time())
    fill_terrain(map)
    return map
end

-- POPULATION CONTROL
local function fill_popul(popul, cull)
    console_log("Populating...")

    math.randomseed(os.time())
    local creature
    if cull then
        creature = {}
        local j = random(1,cull)
        fill_creature(creature, popul[1], popul[j])
        table.insert(popul, creature)
    end

    local popul_size = conditions.popul_max
    while #popul <= popul_size do
        creature = {}
        if conditions.champ_breed
        and (not cull or random() < 0.05) then
            fill_creature(creature)
        else
            local i = random(1,cull)
            local j
            repeat
                j = random(1,cull)
            until i ~= j
            fill_creature(creature, popul[i], popul[j])
        end
        table.insert(popul, creature)
    end

    console_log("Success populating.\n")
end

local function fitness_sort(a, b)
    return a.fitness > b.fitness
end

-- TURN MANAGEMENT
local function move_creature(creature, n)
    local p = {}
    p.x = creature.pos.x + adjacent[n].x
    p.z = creature.pos.z + adjacent[n].z
    if valid(p) then p.y = terrain[p.x][p.z][1] end
    if p.y and p.y - creature.pos.y < creature.climb then
        creature.pos = p
        if creature.pos.y - p.y > creature.fear_height then
            creature.hp = creature.hp - ceil(creature.pos.y - p.y)
        end
    end
end

local function manage_turn(creature, map, turn)
    creature.lifespan = turn
    local last = manage_creature(creature, map, conditions.dtime)

    console_log(conditions.turn_log(creature, map, gen_stats, turn))
    conditions.map_checks(creature, map, turn)

    return last
end

local function manage_lifetime(creature, map)
    local last, turn, max_turn = false, 0, conditions.max_turn

    while turn <= max_turn and not last do
        turn = turn + 1
        last = manage_turn(creature, map, turn)
    end

    if conditions.fitness_modifier then
        creature.fitness = creature.fitness * conditions.fitness_modifier(turn, max_turn)
    end
end

local function manage_gen(num)
    gen_stats.num = num

    if not output and not minetest then
        print("-- Start of generation "..num.." --")
    end

    local s = string.rep("-", 42).."\n"..string.rep("-", 42).."\n\n"
    s = s.."*** START OF GEN "..num.." ***".."\n"
    console_log(s)

    local map = gen_terrain()
    terrain = map

    local creature
    for i=1,#popul do
        creature = popul[i]
        creature.id = i

        reset_creature(creature, map)
        manage_lifetime(creature, map)
    end

    table.sort(popul, fitness_sort)

    local champ = popul[1]

    gen_stats.avg = avg(popul, function(creature)
        return creature.fitness
    end)
    gen_stats.peak = champ.fitness

    local i, elite = 1, gen_stats.top[1]
    while elite and i <= top_size do
        if champ.fitness > elite.fitness then
            gen_stats.top[i] = clone_creature(champ)
            elite = nil
        else
            i = i + 1
        end
    end
    if not gen_stats.top[i] then
        gen_stats.top[i] = clone_creature(champ)
    end

    if champ.fitness < conditions.fitness_threshold then
        popul = {}
        fill_popul(popul)
    else
        if champ.fitness > gen_stats.max_peak then
            gen_stats.gen = num
            gen_stats.max_peak = champ.fitness
            conditions.custom_stats(champ, gen_stats)
        end

        local cull = ceil(#popul*conditions.cull_threshold)
        for i=cull,#popul do
            popul[i] = nil
        end

        fill_popul(popul, cull)
    end

    console_log(conditions.gen_log(num, gen_stats))
end

-- TRAINING PROCESS
local function train(options, offset)
    conditions = {
        dims = options.dims or {x=40, y=40, z=40},
        output = options.output or false,
        gens = options.gens or 100,
        dtime = options.dtime or 1,
        turn_log = options.turn_log or function() end,
        gen_log = options.gen_log or function() end,
        training_log = options.training_log or function() end,
        map_checks = options.map_checks or function() end,
        fitness_threshold = options.fitness_threshold or 0,
        cull_threshold = options.cull_threshold or 3/5,
        champ_breed = options.champ_breed or true,
        max_turn = options.max_turn or 200,
        fitness_modifier = options.fitness_modifier,
        custom_stats = options.custom_stats,
        popul_max = options.popul_max or 150,
    }

    gen_stats = options.gen_stats or {}
    gen_stats.num = 0
    gen_stats.avg = 0
    gen_stats.peak = 0
    gen_stats.max_peak = 0
    gen_stats.top = {}

    popul = options.popul or {}
    fill_cell = options.fill_cell
    fill_creature = options.fill_creature
    reset_creature = options.reset_creature
    manage_creature = options.manage_creature
    clone_creature = options.clone or function(creature)
        return creature
    end
    top_size = options.top_size > 0 and options.top_size or 1

    fill_popul(popul)

    local last_gen = ofsset or 0

    for i=1,conditions.gens do
        manage_gen(i + last_gen)
    end
    last_gen = gen_stats.num
    console_log(conditions.training_log(gen_stats))

    if options.to_file then
        for i=1,top_size do
            s = popul[i]:to_lines()
            file = io.open(options.to_file..".txt", "a")
            for _,line in ipairs(s) do
                file:write(line, "\n")
            end
            file:close()
        end
    end

    return popul, last_gen
end

trainer.train = train
trainer.valid = valid
trainer.console_log = console_log
trainer.fitness_sort = fitness_sort
trainer.move_creature = move_creature

adaptive_ai.trainer = trainer