'''
Problem 4.4

Consider an M=Ek=1=N = 10 queue with lmbda = 60/30 = 2 per hour and mu = 60/20 =
3 per hour. Implement the so-called embedded Markov chain method in Python.
First assume k = 1 (exponential) and then compare the result for the steady state probabilities for
k = 3.

'''
import numpy as np
import MarkovLib as ml
def ndx(i,j): # Index function to get index in A for combination main state i, and sub state j
    if i == 0:
        return 0
    else:    
        return (i-1)*k + j +1
mu = 3  # Total transition rate between main states
lmbda = 2

for k in [1,3]: # Repeat for k = 1 and 3
    N = 10 # Number of places in the system
    Dim = N*k+1
    A = np.zeros([Dim, Dim])
    A[ndx(0,0),ndx(1,0)] = lmbda
    for i in range(1,N):
        for j in range(k):
            A[ndx(i,j),ndx(i+1,j)] = lmbda
            
    for i in range(1,N+1):
        for j in range(k-1):
            A[ndx(i,j),ndx(i,j+1)] = k * mu
        A[ndx(i,k-1),ndx(i-1,0)] = k * mu    
    ml.fixA(A)
    P=ml.P_asymptotic(A)
    print("\nResult for k=",k)
    for i in range(N+1):
        if i == 0:
            pp=P[ndx(0,0)]
        else:
            pp=0
            for j in range(k):
                pp += P[ndx(i,j)]
        print(f'P_{i}:','{:.3E}'.format(pp))