--@name NeuroEvolution Experiments
--@author Szymekk
--@shared

if CLIENT then return end

--------------------------------------
-- Neuron class
--------------------------------------

local neuron = { }
neuron.__index = neuron

function neuron:init()
    self.inputs = { }
    self.weights = { }

    self.value = nil
end

function neuron:activationFunction(x)
    return 1 / (1 + math.exp(-x)) * 2 - 1
    --return x
end

function neuron:addInput(n)
    self.inputs[#self.inputs + 1] = n
    self.weights[#self.inputs] = math.random() * 2 - 1
end

function neuron:getOutput()
    if not self.value and #self.inputs > 0 then
        local sum = 0

        for k, v in pairs(self.inputs) do
            sum = sum + v:getOutput() * self.weights[k]
        end

        self.value = neuron:activationFunction(sum)
    end

    return self.value
end

function neuron.new()
    local tbl = { }
    setmetatable(tbl, neuron)
    tbl:init()

    return tbl
end

--------------------------------------
-- Neural network class
--------------------------------------

local network = { }
network.__index = network

function network:init()
    self.inputs = { }
    self.outputs = { }
    self.hiddenLayers = { }
    self.allNeurons = { }
end

function network:addInputs(count)
    for i = 1, count do
        local n = neuron.new()
        self.inputs[i] = n
        self.allNeurons[#self.allNeurons + 1] = n
    end
end

function network:addHiddenLayer(count)
    local layer = { }

    for i = 1, count do
        local n = neuron.new()
        layer[i] = n
        self.allNeurons[#self.allNeurons + 1] = n
    end

    self.hiddenLayers[#self.hiddenLayers + 1] = layer

    return layer
end

function network:addOutputs(count)
    for i = 1, count do
        local n = neuron.new()
        self.outputs[i] = n
        self.allNeurons[#self.allNeurons + 1] = n
    end
end

function network:setInputs(tbl)
    for k, v in pairs(self.inputs) do
        v.value = tbl[k]
    end
end

function network:resetValues()
    for k, v in pairs(self.allNeurons) do
        v.value = nil
    end
end

function network:connectNeurons()
    local allLayers = { self.inputs }
    for k, v in pairs(self.hiddenLayers) do
        allLayers[#allLayers + 1] = v
    end
    allLayers[#allLayers + 1] = self.outputs

    for i = 2, #allLayers do
        for k1, v1 in pairs(allLayers[i - 1]) do
            for k2, v2 in pairs(allLayers[i]) do
                v2:addInput(v1)
            end
        end
    end
end

function network:run(tbl)
    self:resetValues()
    self:setInputs(tbl)

    local outputValues = { }

    for k, v in pairs(self.outputs) do
        outputValues[k] = v:getOutput()
    end

    return outputValues
end

function network:dumpWeights()
    local weights = { }

    for k, v in pairs(self.allNeurons) do
        weights[k] = { }

        for k2, v2 in pairs(v.weights) do
            weights[k][k2] = v2
        end
    end

    return weights
end

function network:loadWeights(weights)
    for k, v in pairs(weights) do
        for k2, v2 in pairs(v) do
            self.allNeurons[k].weights[k2] = v2
        end
    end
end

function network.new()
    local tbl = { }
    setmetatable(tbl, network)
    tbl:init()

    return tbl
end

--------------------------------------
-- Evolution class
--------------------------------------

local evolution = { }
evolution.__index = evolution

function evolution:init(net)
    self.net = net

    self.netsPerGeneration = 100
    self.maxGenerations = 40

    self.mutationNeurons = 5
    self.mutationRate = 0.01
    self.maxWeight = 5
end

function evolution:start()
    return coroutine.create(function()
        self:yieldedEvolve()
    end)
end

function evolution:yieldedEvolve()
    self:onStart()

    local generation = 1

    local lastBest
    local mutationRateFactor = 1

    while true do
        print("Generation: " .. generation)

        self:onGenerationStart(generation)

        local nets = { }
        local weights = self.net:dumpWeights()

        for i = 1, self.netsPerGeneration do
            local updatedWeights = self:updateWeights(weights, mutationRateFactor)
            self.net:loadWeights(updatedWeights)

            local fitness = self:onUpdate(self.net)

            nets[i] = { fitness, updatedWeights }

            coroutine.yield()
        end

        coroutine.yield()

        local best = nets[1]
        for k, v in pairs(nets) do
            if v[1] < best[1] then
                best = v
            end
        end

        print("Best fitness: ", best[1])

        if lastBest and lastBest[1] * 1 > best[1] then
            print("Not better than before")
            self.net:loadWeights(lastBest[2])
            mutationRateFactor = mutationRateFactor + 0.5
            continue
        end

        self.net:loadWeights(best[2])

        lastBest = best

        mutationRateFactor = 1

        self:onGenerationEnd(generation)

        if generation >= self.maxGenerations then
            break
        end

        generation = generation + 1
    end

    self:onEnd()
end

-- Evolution begins, all preparation is done here
function evolution:onStart() end

-- Update function of current step, should return current fitness
function evolution:onUpdate(net) end

-- Before training individuals
function evolution:onGenerationStart(n) end

-- After best individual has been chosen
function evolution:onGenerationEnd(n) end

-- The network has evolved
function evolution:onEnd() end

function evolution:updateWeights(weights, factor)
    local updated = { }

    for k, v in pairs(weights) do
        updated[k] = { }

        for k2, v2 in pairs(v) do
            updated[k][k2] = v2
        end
    end

    for i = 1, self.mutationNeurons do
        local k = math.random(1, #updated)

        if #updated[k] == 0 then
            i = i - 1
            continue
        end

        local l = math.random(1, #updated[k])
        updated[k][l] = math.clamp(updated[k][l] + (math.random() * 2 - 1) * self.mutationRate * factor, -self.maxWeight, self.maxWeight)
    end

    return updated
end

function evolution.new(net)
    local tbl = { }
    setmetatable(tbl, evolution)
    tbl:init(net)

    return tbl
end

--------------------------------------
-- Action
--------------------------------------

chip():setMaterial("models/shiny")

local net = network.new()
net:addInputs(3)
net:addHiddenLayer(3)
--net:addHiddenLayer(3)
net:addOutputs(1)
net:connectNeurons()

local evo = evolution.new(net)

evo.netsPerGeneration = 10
evo.maxGenerations = 5000

evo.mutationNeurons = 8
evo.mutationRate = 0.2
evo.maxWeight = 5

local evoCoroutine = evo:start()

local allPower = true

function evo:onStart()
    print("Starting evolution")
end

function scaleInput(a)
    return a * 0.1
end

function scaleOutput(a)
    return a * 10
end


function evo:onGenerationStart()
    a, b = math.random(1, 10), math.random(1, 10)
end

startPos = chip():getPos() + Vector(0, 0, 15)

color = Color(math.random(0, 255), math.random(0, 255), math.random(0, 255))

baseHolo = holograms.create(startPos, Angle(0, 0, 0), "models/sprops/cuboids/height12/size_2/cube_24x18x12.mdl", Vector(1, 1, 1))
baseHolo:setColor(color)
traceHolo1 = holograms.create(startPos, Angle(0, 0, 0), "models/sprops/geometry/sphere_6.mdl", Vector(1, 1, 1))
traceHolo2 = holograms.create(startPos, Angle(0, 0, 0), "models/sprops/geometry/sphere_6.mdl", Vector(1, 1, 1))
traceHolo3 = holograms.create(startPos, Angle(0, 0, 0), "models/sprops/geometry/sphere_6.mdl", Vector(1, 1, 1))

local bestFitness = 0

function evo:onUpdate(net)
    allPower = false

    local position = startPos
    local angle = math.pi / 2

    while true do
        position = position + Vector(math.cos(angle), math.sin(angle)) * 2

        baseHolo:setPos(position)
        baseHolo:setAngles(Angle(0, angle / math.pi * 180 + 90, 0))

        local tr1 = trace.trace(position, position + Vector(math.cos(angle), math.sin(angle)) * 100)
        local tr2 = trace.trace(position, position + Vector(math.cos(angle + math.pi / 4), math.sin(angle + math.pi / 4)) * 100)
        local tr3 = trace.trace(position, position + Vector(math.cos(angle - math.pi / 4), math.sin(angle - math.pi / 4)) * 100)

        traceHolo1:setPos(tr1.HitPos)
        traceHolo2:setPos(tr2.HitPos)
        traceHolo3:setPos(tr3.HitPos)

        if tr1.Fraction < 0.01 or tr2.Fraction < 0.01 or tr3.Fraction < 0.01 then
            break
        end

        local rawOutput = net:run({ scaleInput(tr1.Fraction), scaleInput(tr2.Fraction), scaleInput(tr3.Fraction) })[1]
        local scaledOutput = scaleOutput(rawOutput)

        --[[
        if scaledOutput < -0.1 then
            angle = angle + 0.02
        elseif scaledOutput > 0.1 then
            angle = angle - 0.02
        end
        ]]

        angle = angle + scaledOutput * 0.1

        coroutine.yield()
    end

    allPower = true

    local fitness = (position - startPos).y - (position - startPos).x / 2

    if fitness > bestFitness then
        baseHolo:emitSound("resource/warning.wav")
        bestFitness = fitness
    end

    return fitness
end

function evo:onGenerationEnd()
    print("Generation ended")
end

function evo:onEnd()

end

hook.add("think", "", function()
    if allPower then
        while coroutine.status(evoCoroutine) ~= "dead" and quotaUsed() / quotaMax() < 0.1 do
            coroutine.resume(evoCoroutine)
        end
    else
        if coroutine.status(evoCoroutine) ~= "dead" and quotaUsed() / quotaMax() < 0.1 then
            coroutine.resume(evoCoroutine)
        end
    end
end)

--[[
chip():setMaterial("models/shiny")

local net = network.new()
net:addInputs(6)
net:addHiddenLayer(3)
net:addOutputs(3)
net:connectNeurons()

local evo = evolution.new(net)

evo.netsPerGeneration = 15
evo.maxGenerations = 5000

evo.mutationNeurons = 8
evo.mutationRate = 0.05
evo.maxWeight = 5

local evoCoroutine = evo:start()

local allPower = true

function evo:onStart()
    print("Starting evolution")
end

function scaleInput(a)
    return a * 0.001
end

function scaleOutput(a)
    return a * 1000
end

local a, b

function evo:onGenerationStart()
    a, b = math.random(1, 10), math.random(1, 10)
end

startPos = chip():getPos() + Vector(0, 0, 2)
desiredPos = startPos + Vector(50, 30, 60)

holo = holograms.create(desiredPos, Angle(0, 0, 0), "models/sprops/geometry/sphere_9.mdl", Vector(1, 1, 1))

function evo:onUpdate(net)
    allPower = false

    local fitness = 0
    local nextEndTime = timer.realtime() + 2

    chip():setPos(startPos)
    chip():setFrozen(false)
    chip():setVelocity(Vector(0, 0, 0))
    chip():enableGravity(false)

    while timer.realtime() < nextEndTime do
        fitness = fitness + ((desiredPos - chip():getPos()):getLength() ^ 2) * timer.frametime()

        local relPos = chip():getPos() - desiredPos
        local vel = chip():getVelocity()
        local rawOutput = net:run({ scaleInput(relPos.x), scaleInput(relPos.y), scaleInput(relPos.z), scaleInput(vel.x), scaleInput(vel.y), scaleInput(vel.z) })
        local scaledOutput = Vector(scaleOutput(rawOutput[1]), scaleOutput(rawOutput[2]), scaleOutput(rawOutput[3]))

        chip():applyForceCenter(scaledOutput)

        coroutine.yield()
    end

    allPower = true

    return fitness
end

function evo:onGenerationEnd()
    print("Generation ended")
end

function evo:onEnd()

end

hook.add("think", "", function()
    if allPower then
        while coroutine.status(evoCoroutine) ~= "dead" and quotaUsed() / quotaMax() < 0.6 do
            coroutine.resume(evoCoroutine)
        end
    else
        if coroutine.status(evoCoroutine) ~= "dead" and quotaUsed() / quotaMax() < 0.6 then
            coroutine.resume(evoCoroutine)
        end
    end
end)
]]