import sys
import gym
import numpy as np
from collections import defaultdict

# Init gym blackjack env
env = gym.make('Blackjack-v0')

# Get space & action sizes
print(env.observation_space)
print(env.action_space)

# Run a x episodes and chose random actions
def runRandom(episodes_count):
    for i in range(episodes_count):
        print("> Episode",i)
        print("-----------")
        
        observation = env.reset()
        print("+Initial state:",observation)
        while True:
            action = env.action_space.sample()
            state, reward, done, info = env.step(action)
            print("+new step:",state,action,reward,done,info)
            if done:
                print("+End of Game: ")
                print("+You won :) ...",reward) if reward > 0 else print("-You lost :( ...",reward)
                break

# Run an episode with the following policy:
# - [.2,.8] if state <= 18
# - [.8,.2] if state >  18
def generate_episode_from_limit_stochastic(env):
    episode = []
    state = env.reset()
    while True:
        p_sa = [.8,.2] if state[0] > 18 else [.2,.8]
        action = np.random.choice(np.arange(2), p=p_sa)
        next_state, reward, done, _ = env.step(action)
        episode.append((state, action, reward))
        state = next_state
        if done:
            break
    return episode

# Support function to compute the average iteratively
def averageIterative(x_n:float, n:int, x_n_1:float)->float:
    return x_n/n + (n-1)/n * x_n_1

def mc_prediction_q(env, num_episodes, gamma=1.0):
    N=defaultdict(lambda:np.zeros(env.action_space.n))
    Q=defaultdict(lambda:np.zeros(env.action_space.n))
    
    for i in range(num_episodes):
        # monitor progress
        if (i+1) % 1000 == 0:
            print("\rEpisode {}/{}.".format(i+1, num_episodes), end="")
            sys.stdout.flush()
        
        # Generate episode
        episode = generate_episode_from_limit_stochastic(env)
        states, actions, rewards = zip(*episode)
        
        for i in range(len(states)):
            N[states[i]][actions[i]] += 1
            Q[states[i]][actions[i]] = averageIterative(
                gamma**i * rewards[i],
                N[states[i]][actions[i]],
                Q[states[i]][actions[i]]
            )
    return Q
#episode = generate_episode_from_limit_stochastic(env)
#print(episode)

Q=mc_prediction_q(env, 500000)
print(Q)
env.close()
import sys
import gym
import numpy as np
from collections import defaultdict

# Init gym blackjack env
env = gym.make('Blackjack-v0')

# Get space & action sizes
print(env.observation_space)
print(env.action_space)

# Run a x episodes and chose random actions
def runRandom(episodes_count):
    for i in range(episodes_count):
        print("> Episode",i)
        print("-----------")
        
        observation = env.reset()
        print("+Initial state:",observation)
        while True:
            action = env.action_space.sample()
            state, reward, done, info = env.step(action)
            print("+new step:",state,action,reward,done,info)
            if done:
                print("+End of Game: ")
                print("+You won :) ...",reward) if reward > 0 else print("-You lost :( ...",reward)
                break

# Run an episode with the following policy:
# - [.2,.8] if state <= 18
# - [.8,.2] if state >  18
def generate_episode_from_limit_stochastic(env):
    episode = []
    state = env.reset()
    while True:
        p_sa = [.8,.2] if state[0] > 18 else [.2,.8]
        action = np.random.choice(np.arange(2), p=p_sa)
        next_state, reward, done, _ = env.step(action)
        episode.append((state, action, reward))
        state = next_state
        if done:
            break
    return episode

# Support function to compute the average iteratively
def averageIterative(x_n:float, n:int, x_n_1:float)->float:
    return x_n/n + (n-1)/n * x_n_1

def mc_prediction_q(env, num_episodes, gamma=1.0):
    N=defaultdict(lambda:np.zeros(env.action_space.n))
    Q=defaultdict(lambda:np.zeros(env.action_space.n))
    
    for i in range(num_episodes):
        # monitor progress
        if (i+1) % 1000 == 0:
            print("\rEpisode {}/{}.".format(i+1, num_episodes), end="")
            sys.stdout.flush()
        
        # Generate episode
        episode = generate_episode_from_limit_stochastic(env)
        states, actions, rewards = zip(*episode)
        
        for i in range(len(states)):
            N[states[i]][actions[i]] += 1
            Q[states[i]][actions[i]] = averageIterative(
                gamma**i * rewards[i],
                N[states[i]][actions[i]],
                Q[states[i]][actions[i]]
            )
    return Q
#episode = generate_episode_from_limit_stochastic(env)
#print(episode)

Q=mc_prediction_q(env, 500000)
print(Q)
env.close()
# Init gym blackjack env # Get space & action sizes # Run a x episodes and chose random actions # Run an episode with the following policy: # - [.2,.8] if state <= 18 # - [.8,.2] if state >  18 # Support function to compute the average iteratively # monitor progress # Generate episode #episode = generate_episode_from_limit_stochastic(env) #print(episode)