dofile("./trainer.lua")
dofile("../kobold/kobold.lua")

local trainer = adaptive_ai.trainer
local helper = adaptive_ai.helper
local calc = adaptive_ai.calc
local kobolds = adaptive_ai.kobolds
local file_tag = kobolds.trained_tag

local random = math.random
--local mod = math.fmod or math.mod
local move_creature = trainer.move_creature
local truncate = calc.truncate
local distance = calc.dist_euclidean
local pop = helper.pop
local valid = trainer.valid

local adjacent = helper.adjacent
local turn_max = 300
local popul_size = 100

local function fill_cell_kobold(cell, y)
    if random() <= 0.01 then
        -- Putting impassable obstacle
        cell[1] = nil
    else
        if random() <= 0.02 then
            -- Planting mushroom
            cell[2] = true
        else
            cell[2] = false
        end
    end
end

local function fill_kobold(kobold, mother, father)
    kobold.fear_height = kobolds.stats.fear_height
    kobold.pos = {x=random(1,40), z=random(1,40)}

    kobolds.setup(kobold, mother, father)

    kobold.fitness = 0
    trainer.console_log("Created new kobold: "..kobold.name)
end

local function reset_kobold(kobold, map)
    kobold.food = {}
    kobold.fitness = 0
    kobold.hp = kobolds.stats.hp_max
    kobold.hunger = kobolds.stats.hunger_max


    kobold.pos.y = map[kobold.pos.x][kobold.pos.z][1]
    while not kobold.pos.y do
        kobold.pos = {x=random(1,40), z=random(1,40)}
        kobold.pos.y = map[kobold.pos.x][kobold.pos.z][1]
    end
end

local kobold_decisions = {
    function(self, dtime) -- idle
    end,
    function(self, dtime) -- wander
        local n = random(1,#adjacent)
        move_creature(self, n)
        self.hunger = self.hunger - 1
    end,
    function(self, dtime) -- eat
        if self.food and #self.food > 0 then
            local satiate = pop(self.food)
            self:eat(satiate)
        end
    end,
    function(self, dtime) -- breed
        self.fitness = self.fitness + 1
        self.hunger = self.hunger - 10
    end,
    function(self, dtime) -- move North
        move_creature(self, 5)
    end,
    function(self, dtime) -- move South
        move_creature(self, 4)
    end,
    function(self, dtime) -- move East
        move_creature(self, 7)
    end,
    function(self, dtime) -- move West
        move_creature(self, 2)
    end,
    function(self, dtime) -- move NE
        move_creature(self, 8)
    end,
    function(self, dtime) -- move SE
        move_creature(self, 6)
    end,
    function(self, dtime) -- move NW
        move_creature(self, 3)
    end,
    function(self, dtime) -- move SW
        move_creature(self, 1)
    end,
}

local function manage_kobold(player, map, dtime)
    if player.hp > 0 then
        local neighbors = {}
        local p, pos = {}, player.pos
        for i=1,#adjacent do
            p.x = pos.x + adjacent[i].x
            p.z = pos.z + adjacent[i].z
            if valid(p) and map[p.x][p.z][1] then
                p.y = map[p.x][p.z][1] - player.pos.y
            else
                p.y = pos.y + player.climb + 1
            end
            neighbors[i] = {x=p.x, y=p.y, z=p.z}
            --print("Neighbor "..i.." = "..p.y)
        end

        player.hunger = player.hunger - dtime
        local decision, gene = player.decide(player, neighbors)
        if decision and decision ~= -1 then
            player.action = decision
            player.last_gene = gene
            kobold_decisions[decision](player, dtime)
        else
            if not player.action then player.action = 1 end
        end

        pos = player.pos

        if map[pos.x][pos.z][2] then
            table.insert(player.food, random(2,5))
            map[pos.x][pos.z][2] = false
            if player.closest_food
            and player.closest_food.x == pos.x
            and player.closest_food.z == pos.z then
                player.closest_food = nil
            end
        end

        if player.hunger < 0 then
            player.hunger = 0
            player.hp = player.hp - 1
        end
        player.fitness = player.fitness + player:score()
        return false
    else
        return true
    end
end

local function print_terrain(map, player)
    local s = string.rep("#",42).."\n"
    local line, cell
    local p = player.pos

    for i=1,40 do
        line = "#"
        for j=1,40 do
            if p.x == i and p.z == j then
                cell = "i"
            elseif map[i][j][1] == nil then
                cell = "O"
            elseif map[i][j][1] == p.y and map[i][j][2] == true then
                cell = "T"
            elseif map[i][j][1] == p.y then
                cell = "▒"
            elseif map[i][j][1] - p.y == -1 then
                cell = "▓"
            elseif map[i][j][1] - p.y == 1 then
                cell = "░"
            elseif map[i][j][1] - p.y < -1 then
                cell = "█"
            elseif map[i][j][1] - p.y > 1 then
                cell = " "
            else
                cell = "?"
            end
            line = line..cell
        end
        s = s..line.."#\n"
    end
    s = s..string.rep("#",42)
    return s
end

local function final_fitness(player, turn, max_turn)
    local c1 = turn
    local c2 = player.fitness
    local m = 3 -- * max_turn
    return (m*c1 + c2)/100
end

local function turn_log_kobold(player, map, gen_stats, turn)
    local s = "\n\nBest creature: "..gen_stats.champ.."\tBest creature fitness: "..truncate(gen_stats.max_peak).."\n"
    s = s.."Best creature lifespan: "..(gen_stats.lifespan or "None").."\tBest creature gen: "..(gen_stats.gen or "None").."\n"
    s = s.."Last gen avg: "..truncate(gen_stats.avg).."\t\tLast gen peak: "..truncate(gen_stats.peak).."\n\n"
    s = s.."Current gen: "..gen_stats.num.."\t\tCurrent creature: "..player.id.."\n"
    s = s.."Turn: "..(turn < 10 and "0" or "")..turn.."\t\tCreature name: "..player.name.."\n"
    s = s.."Action: "..kobolds.actions[player.action].."\t\tFocus: "..player.focus.."\n"
    s = s.."Gene: "..(player.last_gene or "None").."\t\tFitness: "..truncate(final_fitness(player, turn, turn_max)).."\n\n"

    s = s..print_terrain(map, player)
    return s, 0.008
end

local function gen_log_kobold(num, gen_stats)
    local s = "\n\nBest creature: "..gen_stats.champ.."\tBest creature fitness: "..truncate(gen_stats.max_peak).."\n"
    s = s.."Best creature lifespan: "..(gen_stats.lifespan or "None").."\tBest creature gen: "..(gen_stats.gen or "None").."\n"
    s = s.."Last gen average: "..truncate(gen_stats.avg).."\t\tLast gen peak: "..truncate(gen_stats.peak).."\n\n"
end

local function map_checks_kobold(player, map)
    if random() <= 0.05 then
        map[random(1,40)][random(1,40)][2] = true
    end

    for i=1,40 do
        for j=1,40 do
            if map[i][j][1] and map[i][j][2] == true then
                local pos = {x=i, y=map[i][j][1], z=j}
                if not pos.y then error("No pos.y in "..i..","..j) end
                if not player.pos.y then error("No player.pos.y") end
                if player.closest_food and not player.closest_food.y then error("No player.closest_food.y") end
                if not player.closest_food
                or distance(player.pos, pos) < distance(player.pos, player.closest_food) then
                    player.closest_food = pos--{x=pos.x, y=pos.y, z=pos.z}
                end
            end
        end
    end
end

local function gen_stats_update(champ, gen_stats)
    gen_stats.champ = champ.name
    gen_stats.peak = champ.fitness
    gen_stats.lifespan = champ.lifespan
end

local function global_log(gen_stats)
    local s = string.rep("-", 42).."\n"..string.rep("-", 42).."\n\n"
    s = s.."Last champion: "..gen_stats.champ.."\n"
    s = s.."From generation "..gen_stats.gen.."\tTotal generations: "..gen_stats.num.."\n"
    s = s.."Fitnes score: "..gen_stats.max_peak.."\tLifespan: "..(gen_stats.lifespan or "None").."\n"
    s = s.."\n"..string.rep("-", 42).."\n"..string.rep("-", 42).."\n"

    return s
end

local function train()
    local input = false
    local option = input and 0 or 300
    local offset = 0
    local population = {}
    repeat
        local gens_option = tonumber(option)
        if gens_option and gens_option > 0 then
            local elapsed = os.clock()
            local stats = {
                champ = "None",
                lifespan = 0,
                gen = 0,
            }

            population, offset = trainer.train({
                dims = {x=40, y=10, z=40},
                output = input,
                gens = gens_option,
                fill_cell = fill_cell_kobold,
                fill_creature = fill_kobold,
                reset_creature = reset_kobold,
                decision_managers = kobold_decisions,
                manage_creature = manage_kobold,
                clone = kobolds.clone,
                map_checks = map_checks_kobold,
                turn_log = turn_log_kobold,
                gen_log = gen_log_kobold,
                training_log = global_log,
                fitness_threshold = 40,
                fitness_func = final_fitness,
                max_turn = turn_max,
                popul = population,
                custom_stats = gen_stats_update,
                gen_stats = stats,
                to_file = file_tag,
                top_size = 4,
            }, offset)

            if not minetest then
                elapsed = os.clock() - elapsed

                print(global_log(stats))
                print()
                print("Elapsed time: ", elapsed)
            end
        end

        if minetest or not input then
            option = "q"
        else
            print("Please enter number of gens to simulate:")
            option = io.read()
        end
    until option == "q" or option == "Q"
end

train()