import numpy as np
import random

# 输入3*3数组，输出元组
def get_tuple(numpy_array):
    return tuple(map(tuple, numpy_array))

def print_board(board):
    for i in range(3):
        print('-------------')
        out = '| '
        for j in range(3):
            if board[i, j] == 1:
                token = '*'
            elif board[i, j] == -1:
                token = 'x'
            else:
                token = '0'
            out += token + ' | '
        print(out)
    print('-------------')

class Agent_my():
    # 参数：玩家索引号，玩家训练参数（Epsilon、Learning_rate）
    # 玩家需要有游戏策略，策略表，after_state,q,r
    def __init__(self,ooxx_index,Epsilon=0.1,Learning_rate=0.1):
        self.index = ooxx_index
        self.epsilon = Epsilon
        self.alpha = Learning_rate
        self.q_values = {}  # Q值表，创建字典
        self.outcomes = []
        self.isepsilon = [] # 记录该步骤是否探索了
        # 字典的键是棋盘可能的状态，值是状态对应的估计收益值

    def reset(self):
        self.outcomes =[]
        self.isepsilon = []

    # 输入当前棋盘状态
    # 选择行动，根据Q值表选择，搜索当前状态下的所有可能的键，
    # 返回下一棋局的状态，即可
    def choose_action(self, current_board):
        valid_moves=[]
        # 先计算所有的outcome
        for row in range(3):
            for col in range(3):
                outc = current_board.copy()
                if outc[row, col] == 0:
                    # 返回outcome
                    outc[row, col] = self.index
                    valid_moves.append(get_tuple(outc))

        if random.uniform(0, 1) < self.epsilon:
            self.isepsilon.append(0)
            # 随机选择动作
            next_board = random.choice(valid_moves)
        else:
            self.isepsilon.append(1)
            # 选择具有最大Q值的动作
            actions ={}
            max_q = 0
            for action in valid_moves:
                temp_q = self.q_values.setdefault(action, 0)
                actions[action] = temp_q
                if temp_q > max_q:
                    max_q = temp_q
            # 得到了最大的q值和actions列表
            # keys_with_value = [key for key, value in my_dict.items() if value == target_value]
            best_actions = [action for action, q_value in actions.items() if q_value == max_q]
            next_board = random.choice(best_actions)

        # 将选取的outcome加入列表
        self.outcomes.append(next_board)

        return np.array(next_board)


    # 更新Q值
    # 输入是reward，current_board
    # 和棋奖励是0.5，赢棋奖励是1，输棋奖励0
    def update_q_value(self, reward):
        # 先单独设置T时刻
        self.q_values[self.outcomes[-1]] = reward

        # 然后设置T-1,T-2...
        for i in reversed(range(len(self.outcomes) - 1)):
            current_t = self.outcomes[i]
            last_T = self.outcomes[i+1]

            # 下面就是程序的关键，计算公式，蒙特卡洛法的计算公式
            # 这里setdefault是为了避免某个键尚不存在
            self.q_values[current_t] = self.q_values.setdefault(current_t,0) + self.alpha * self.isepsilon[i] * (self.q_values[last_T] - self.q_values.setdefault(current_t,0))

# 返回true，false
def board_over(board,player_index):
    # Check rows, columns, and diagonals
    for i in range(3):
        # all()函数用于判断整个数组中的元素的值是否全部满足条件，如果满足条件返回True，否则返回False
        if np.all(board[i, :] == player_index) or \
            np.all(board[:, i] == player_index):
            return True
    if np.all(np.diag(board) == player_index) or \
        np.all(np.diag(np.fliplr(board)) == player_index):
        return True
    return False

class Environment_my():
    def __init__(self):
        self.board = np.zeros((3,3))
        self.winner = None
        self.r1 = 0
        self.r2 = 0
        # 记录胜率
        self.win1 =0
        self.win2 = 0
        self.win0 =0

    def reset(self):
        self.board = np.zeros((3, 3))
        self.winner = None
        self.r1 = 0
        self.r2 = 0

    # 检查游戏是否结束或是和棋
    # 返回true 结束，false未结束,返回1，1号玩家胜利，-1，2号玩家胜利
    def is_game_over(self):
        if board_over(self.board,1) == True:
            self.winner = 1
            self.r1 = 1
            self.win1 +=1
            self.r2 = 0
            return True
        if board_over(self.board,-1) == True:
            self.winner = -1
            self.win2 +=1
            self.r1 = 0
            self.r2 = 1
            return True
        if (np.all(self.board != 0) and self.winner==None):
            self.win0 +=1
            self.r1 = 0.5
            self.r2 = 0.5
            return True
        return False

# human interface
# input a number to put a chessman
# | q | w | e |
# | a | s | d |
# | z | x | c |
class HumanPlayer:
    def __init__(self, **kwargs):
        self.symbol = None
        self.keys = ['q', 'w', 'e', 'a', 's', 'd', 'z', 'x', 'c']
        self.state = None

    def reset(self):
        pass
    def choose_action(self, current_board):
        print_board(current_board)
        key = input("Input your position:")
        data = self.keys.index(key)
        i = data // 3
        j = data % 3
        next_board = current_board.copy()
        next_board[i,j] = -1
        return next_board

def human_test(ai_agent,human_player):
    env = Environment_my()
    while True:
        ai_agent.reset()
        env.reset()
        while True:
            env.board = agent1.choose_action(env.board)
            # print(env.board)
            # 检查游戏是否结束
            if env.is_game_over() == True:
                break
            env.board = human_player.choose_action(env.board)
            if env.is_game_over() == True:
                break
        if env.winner == 1:
            print("You lose!")
        elif env.winner == -1:
            print("You win!")
        else:
            print("It is a tie!")


def test_ai(agent1,agent2,testi=1000):
    agent1.epsilon =0
    agent2.epsilon =0
    env_pk =Environment_my()
    for i in range(testi):
        env_pk.reset()
        agent1.reset()
        agent2.reset()
        while True:
            env_pk.board = agent1.choose_action(env_pk.board)
            if env_pk.is_game_over() == True:
                # agent1.update_q_value(env_pk.r1)
                # agent2.update_q_value(env_pk.r2)
                break
            env_pk.board = agent2.choose_action(env_pk.board)
            if env_pk.is_game_over() == True:
                # agent1.update_q_value(env_pk.r1)
                # agent2.update_q_value(env_pk.r2)
                break
    print(f"Agent 1 Wins: {env_pk.win1 / testi}, Agent 2 Wins: {env_pk.win2 / testi}, Draws: {env_pk.win0 / testi}")


if __name__ == "__main__":
    agent1 = Agent_my(1,0.1,0.1)
    agent2 = Agent_my(-1,0.1,0.1)
    env = Environment_my()

    print("start")
    # 落子
    all_e =100000
    z_e= 1000
    for i in range(all_e):
        env.reset()
        agent1.reset()
        agent2.reset()
        while True:
            env.board = agent1.choose_action(env.board)
            # 检查游戏是否结束
            if env.is_game_over() == True:
                agent1.update_q_value(env.r1)
                agent2.update_q_value(env.r2)
                break
            env.board = agent2.choose_action(env.board)
            if env.is_game_over() == True:
                agent1.update_q_value(env.r1)
                agent2.update_q_value(env.r2)
                break
        if (i + 1) % z_e == 0:
            print(f"Agent 1 Wins: {env.win1 / i}, Agent 2 Wins: {env.win2 / i}, Draws: {env.win0 / i}")
    test_ai(agent1,agent2)


    # 训练完成后，与人类玩家竞赛，看看效果。这里设置了人类玩家后手，机器人先手。
    p2 = HumanPlayer()
    human_test(agent1,p2)