import io
import numpy

# S = {0, 1}
# A = {0, 1}
T = []
R = [] 
for s in range(0, 2):
    Ts = []
    Rs = [] 
    for a in range(0, 2):
        Ta = []
        Ra = [] 
        for sPrime in range(0, 2):
            TsasPrime = 0
            Ta.append(TsasPrime)
            RsasPrime = 0
            Ra.append(RsasPrime)
            
        Ts.append(Ta)
        Rs.append(Ra)

    T.append(Ts)
    R.append(Rs)

print(T)
print(R)

T[0][0][0] = 0.25
T[0][0][1] = 1 - T[0][0][0]
T[0][1][0] = 1
T[0][1][1] = 1 - T[0][1][0]
T[1][0][0] = 0.5
T[1][0][1] = 1 - T[1][0][0]
T[1][1][0] = 1
T[1][1][1] = 1 - T[1][1][0]

R[0][0][0] = 1
R[0][0][1] = 0
R[0][1][0] = 2
R[0][1][1] = 0#Dummy
R[1][0][0] = -1
R[1][0][1] = 3
R[1][1][0] = 0
R[1][1][1] = 0

gamma = 0.5

policy = []
policy.append(0)
policy.append(1)

def getNextState(s, a):
    r = numpy.random.rand()
    if(r < T[s][a][0]):
        return 0
    return 1

def cumulativeReward(startState, horizon):
    state = startState

    cReward = 0
    discount = 1
    for t in range(0, horizon):
        action = policy[state]
        nextState = getNextState(state, action)
        reward = R[state][action][nextState]
        cReward += discount * reward
        discount *= gamma
        state = nextState

    return cReward


numRuns = 10000
total = 0
for run in range(0, numRuns):
    total += cumulativeReward(0, 20)

expectedCumulativeReward = total / numRuns

print("Cumulative reward : " + str(expectedCumulativeReward))