import numpy as np

# 状态空间大小
state_size = 5
# 动作空间大小
action_size = 1
# 学习率
learning_rate = 0.5
# 折扣因子
discount_factor = 0.95
# 探索率
epsilon = 0.1
# 最大迭代次数
max_iterations = 10000

# Q-table初始化
q_table = np.zeros((state_size, action_size))
print(q_table)


# Q-learning算法
def q_learning():
    for iteration in range(max_iterations):
        state = np.random.randint(state_size)  # 随机选择一个初始状态
        done = False
        while not done:
            if np.random.uniform() < epsilon:  # 以epsilon的概率随机选择动作
                action = np.random.randint(action_size)
            else:  # 以1-epsilon的概率选择最优动作
                action = np.argmax(q_table[state, :])
            next_state = get_next_state(state, action)  # 根据当前状态和动作得到下一个状态
            reward = get_reward(next_state)  # 根据下一个状态得到奖励值
            update(state, action, next_state, reward)  # 更新Q-table
            state = next_state  # 将下一个状态作为当前状态继续迭代
            if state == get_final_state():  # 到达终止状态，结束迭代
                done = True
    return q_table


# 根据当前状态和动作得到下一个状态和奖励值
def get_next_state(state, action):
    if action == 0:  # 下一步状态为0或1，根据当前状态决定
        return state if state == 0 else 1
    else:  # 下一步状态为1或0，根据当前状态决定
        return state if state == 1 else 0

    # 根据下一个状态得到奖励值，如果到达终止状态则奖励值为1，否则为0


def get_reward(next_state):
    if next_state == get_final_state():
        return 1  # 到达终止状态，奖励值为1
    else:
        return 0  # 没有到达终止状态，奖励值为0


# 返回终止状态的索引值，这里假设终止状态为0或1中的任意一个
def get_final_state():
    return 0 if np.random.randint(2) == 0 else 1  # 以50%的概率返回0，以50%的概率返回1作为终止状态


# 更新Q-table
def update(state, action, next_state, reward):
    old_value = q_table[state, action]
    next_max = np.max(q_table[next_state, :])
    new_value = old_value + learning_rate * (reward + discount_factor * next_max - old_value)
    q_table[state, action] = new_value


# 执行Q-learning算法并打印Q-table的最大值和最小值
q_table = q_learning()
print("Q-table的最大值:", np.max(q_table))
print("Q-table的最小值:", np.min(q_table))