main.py

import numpy as np

Q = 2  # no of workers
S = 3  # size of queues
T = 2  # total amount of jobs
finish_probabilities = (1/(T+1))*np.ones(T+1)  # probability that job is finished by worker. for simplicity assumed uniform


def get_random_state(Q, S, T):
    s = []
    for _ in range(Q):
        q = np.random.randint(low=0, high=T+1, size=S)
        q = list(q)
        q.sort(reverse=True)
        s.append(q)
    job = np.random.randint(low=1, high=T)
    x = [s, job]
    return x


def g(x, u):
    cost_to_go = np.sum(np.count_nonzero(np.array(x[0])))
    if x[1] != 0 and u == 0:
        cost_to_go = cost_to_go + 5.
    return cost_to_go


def f(x, u, w):
    x_next = x
    curr_job = x[1]

    for worker, queue in enumerate(x[0]):
        if finish_probabilities[worker] < w[worker]:
            queue.remove(queue[0])
            x_next[0][worker] = [*queue, 0]
    print(x_next)
    if u != [0]:
        chosen_queue = x_next[0][u-1]
        x_next[0][u-1][chosen_queue.index(0)] = curr_job
        new_job = np.random.randint(low=0, high=T+1)
        x_next[1] = new_job

    return x_next


def u(x):
    u_state = 0
    if x[1] == 0:
        return [u_state]
    else:
        tmp = [idx+1 for idx, queue in enumerate(x[0]) if np.count_nonzero(queue) < S]
        u_state = [u_state, *tmp]
        return u_state


def Vk(x, u, Vk_next):
    # computes expectation and finds control that acheives minimum
    expectations = []
    for control in u:
        gk = g(x, control)
        expectations.append(gk + Vk_next)
    Vk = np.min(expectations)
    optimal_control = np.argmin(expectations)
    return Vk, optimal_control


def DP_algorithm(V_N, N, final_state):
    costs = []
    costk = V_N
    xk = final_state
    uk = u(xk)
    for k in range(N, 0, -1):
        costk, policy = Vk(xk, uk, costk)
        costs.append(costk)
        print(costk, policy)
        wk = np.random.uniform(low=0, high=1, size=Q)
        print(uk, wk)
        xk = f(xk, policy, wk)
        uk = u(xk)


def test():
    final_state = [[[0, 0, 0], [0, 0, 0]], 2]
        # get_random_state(Q, S, T)
    print(final_state)
    VN, N = 3, 10
    DP_algorithm(VN, N, final_state)


if __name__ == '__main__':
    test()