coeffs_matrix = {}

function generateCoeffs(products)
    data = {}
    for i=0, products:size()-1 do
        local product = products:get(i)
        local inputs = product.cost.inputs
        data[i] = {}
        for j=0, inputs:size()-1 do
            local input = inputs:get(j)
            if data[input.product_id] == nil then
                data[input.product_id] = {}
            end
            local row = data[input.product_id]
            row[product.id] = input.quantity
        end
    end
    return data
end

function jacobiPass(products, coeffs_matrix, values, targets)
    local out = {}
    for i=0, products:size()-1 do
        local coeffs = coeffs_matrix[i]
        local target = targets[i]
        local gross = target
        for j,coeff in pairs(coeffs) do
            gross = gross + values[j]*coeff
        end
        out[i] = gross
    end
    return out
end

function jacobi(products, coeffs_matrix, targets, iterations)
    local values = {}
    for i=0, products:size()-1 do
        values[i] = 0
    end
    for i=1, iterations do
        values = jacobiPass(products, coeffs_matrix, values, targets)
    end
    return values
end

function genTargets(products, inventory, consumed, reserve, padding)
    local data = {}
    for i=0, products:size()-1 do
        local consumption = consumed:get(i) * padding
        local stock = inventory:get(i)
        local target = consumption + math.max(0, reserve - stock)
        data[i] = target
    end
    return data
end

lowAdjust = 3
invLowAdjust = 1/lowAdjust
highAdjust = 1/3
invHighAdjust = 1/highAdjust

function harmony(goal, output)
    local u = (output - goal) / goal
    if u < 0 then
        return -(u*u)*lowAdjust
    else
        return math.sqrt(u)*highAdjust
    end
end

function invHarmony(goal, h)
    local u = 0
    if h < 0 then
        local temp = h * invLowAdjust
        u = -math.sqrt(math.abs(temp))
    else
        local temp = h * invHighAdjust
        u = (temp*temp)
    end
    return (u*goal) + goal
end

function insertOrAdd(data, index, value)
    if(data[index] == nil) then
        data[index] = value
    else
        data[index] = data[index] + value
    end
end

function copyTable(a)
    data = {}
    for i,v in pairs(a) do data[i] = v end
    return data
end

function addTable(a, b)
    data = copyTable(a)
    for i,v in pairs(b) do insertOrAdd(data, i, v) end
    return data
end

function subTable(a, b)
    data = copyTable(a)
    for i,v in pairs(b) do insertOrAdd(data, i, -v) end
    return data
end

function addToInv(inventory, a)
    for i,v in pairs(a) do insertOrAdd(inventory, i, v) end
end

function subFromInv(inventory, a)
    for i,v in pairs(a) do insertOrAdd(inventory, i, -v) end
end

function getInputs(product, output)
    data = {}
    local inputs = product.cost.inputs
    for i=0,inputs:size()-1 do
        local input = inputs:get(i)
        data[input.product_id] = input.quantity * output
    end
    data.labor = product.cost.labor * output
    return data
end

function getProduction(product, assigned)
    local production = assigned.labor / product.cost.labor
    local inputs = product.cost.inputs
    for i=0,inputs:size()-1 do
        local input = inputs:get(i)
        local a = assigned[input.product_id]
        if a == nil then a = 0 end
        local q = a/input.quantity
        production = math.min(production, q)
    end
    return production
end

function removeExcess(inventory, assigned, product, output)
    addToInv(inventory, assigned[product.id])
    local inputs = getInputs(product, output)
    subFromInv(inventory, inputs)
    assigned[product.id] = inputs
end

function lessThan(a, b) return a[2] > b[2] end
function balance(products, required, assigned, inventory, bias)
    local ordered = {}
    local avg = 0
    local n = 0
    for i,v in pairs(assigned) do
        local product = products:get(i)
        local production = getProduction(product, v)
        local h = harmony(required[i], production)
        table.insert(ordered, {product.id, h})
        avg = avg + h
        n = n + 1
    end
    
    table.sort(ordered, lessThan)
    avg = avg / n
    avgb = avg + bias
    
    for i=1,#ordered do
        local pair = ordered[i]
        local id = pair[1]
        local h = pair[2]
        local product = products:get(id)
        local old = assigned[id]
        local goal = required[id]
        if h > avgb then
            local output = invHarmony(goal, avgb)
            removeExcess(inventory, assigned, product, output)
        end
        if h < avgb then
            addToInv(inventory, old)
            local max = getProduction(product, inventory)
            local output = math.min(max, invHarmony(goal, avgb))
            local inputs = getInputs(product, output)
            subFromInv(inventory, inputs)
            assigned[id] = inputs
        end
    end
end

function toTable(userdata)
    local data = {}
    for i=0,userdata:size()-1 do
        data[i] = userdata:get(i)
    end
    return data
end

function assignResources()
    local products = econ.products
    
    local agents = econ.getAgents()
    local tasks = econ.getTasks()
    local inventory = econ.getInventory()
    local consumed = econ.getConsumedByAgents()
    
    local avg_labor = econ.getAvgLabor()
    local inv_avg_labor = 1/avg_labor
    local total_labor = avg_labor * agents:size()
    
    local reserve = agents:size() --arbitrary minimum target stock
    
    local targets = genTargets(products, inventory, consumed, reserve, 1)
    local gross_outputs = jacobi(products, coeffs_matrix, targets, 20)
    
    
    local virtual_inv = toTable(inventory)
    virtual_inv.labor = total_labor
    
    local assigned = {}
    for i=0, products:size()-1 do
        if assigned[i] == nil then assigned[i] = {} end
        assigned[i].labor = 0
    end
    
    local bias = 0.1
    local dec = 1 --change this to increase or decrease B per iteration
    for i=1,32 do
        balance(products, gross_outputs, assigned, virtual_inv, bias)
        bias = bias * dec
    end
    
    local data = {}
    local index = 0
    local agent_index = 0
    for i,v in pairs(assigned) do
        data[index+1] = {
            task_id = i,
            inputs = {},
            agents = {},
        }
        
        local agents_number = math.abs(math.ceil(v.labor/avg_labor))
        agents_number = math.min(agents:size()-agent_index, agents_number)
        v.labor = nil
        
        local itr = 0
        for id,q in pairs(v) do
            data[index+1].inputs[itr+1] = { product_id = id, quantity = q }
            itr = itr + 1
        end
        
        for j=1,agents_number do
            local agent = agents:get(agent_index+(j-1))
            data[index+1].agents[j] = agent
        end
        agent_index = agent_index + agents_number
        index = index + 1
        
        if (agent_index >= agents:size()) then break end
    end
    
    return data
end

function assignWages()
    local agents = econ.getAgents()
    data = {}
    for i=1, agents:size() do
        data[i] = {
            agent_id = i-1,
            wage = econ.getAvgLabor(),
        }
    end
    return data
end

function assignPrices()
    data = {}
    for i=1, econ.products:size() do
        local product = econ.products:get(i-1)
        data[i] = {
            product_id = product.id,
            price = product.cost.total,
        }
    end
    return data
end

function printBar()
    print("---------------------------")
end

function printEnd()
    print(
        "day",
        econ.getDay(),
        "average health",
        econ.getAverageHealth(),
        "average happiness",
        econ.getAverageHappiness()
    )
    print("###########################")
    print("")
    io.flush()
end

function printInventory()
    printTab("day "..econ.getDay().." inventory", econ.getInventory())
end

function printConsumption()
    printTab("day "..econ.getDay().." consumption (agents)", econ.getConsumedByAgents())
end

function printProduction()
    printTab("day "..econ.getDay().." production", econ.getProduced())
    printEnd()
end

function printTab(title, tab)
    print(title)
    printBar()
    for i=0, tab:size()-1 do
        local data = tab:get(i)
        if data > 0 then
            print("product_id", i, "quantity", data)
        end
    end
    printBar()
end

function start()
    local products = econ.products
    coeffs_matrix = generateCoeffs(products)
end