import numpy as np
import tensorflow as tf

NUM_STEPS = 10


def get_service_time(inter_arrivals, hidden_state):
    return np.random.rand()

def get_next_arrival(x, hidden_state):
    return np.random.rand()

def predict_service_time(arrivals, arrivals_state, service_state):
    """
    arrivals: [B, 1, D]
    arrivals_state: [B, H_A]
    service_state: [B, H_S]
    """
    service_time = np.zeros(arrivals.shape[0], 1)  # [B]
    service_state = np.zeros_like(service_state)
    inter_arrivals = arrivals
    condition = True
    a_next = arrivals
    a_state = arrivals_state
    while np.count_nonzero(service_time) < arrivals.shape[0]:
        expected_service_time, service_state = get_service_time(inter_arrivals, service_state)
        a_next, a_state = get_next_arrival(a_next, a_state)

        condition = np.greater(expected_service_time, a_next)
        # sould we add the expected departure time in the calculation
        inter_arrivals = np.concatenate((arrivals, a_next), 1)
        service_time[~condition] = expected_service_time[~condition]
        service_state[~condition] = service_state[~condition]

    return service_time, service_state


arrivals = None
arrivals_state = None
service_state = None
for _ in range(NUM_STEPS):
    service_time, service_state = predict_service_time(
        arrivals[:, -1, None], arrivals_state, service_state)
    arrivals = np.concatenate((arrivals, service_time), 1)
    next_arrival, arrivals_state = get_next_arrival(arrivals, arrivals_state)
    arrivals = np.concatenate((arrivals, next_arrival), 1)
