# -*- coding: utf-8 -*-
# @Time       : 2020/10/1 16:30
# @Author     : Duofeng Wu
# @File       : client.py
# @Description:


import json
import time

from ws4py.client.threadedclient import WebSocketClient

import greedy_rule_v2 as greedy_rule
from state import State
from action import Action
import threading
import copy
import my_utils
import numpy as np
from open_spiel.python import rl_environment
from open_spiel.python.algorithms import policy_gradient
from open_spiel.python.algorithms import dqn

import tensorflow.compat.v1 as tf
import os
import sys
import strong_rule_client as opponent
import example_clients.RandomClient as rdClient
import AI_Code.gdai2 as AI
import my_bucket_action_strong_rule as bucket_rule

os.environ['CUDA_VISIBLE_DEVICES']='0,1'
tf.disable_v2_behavior()
save_dir = 'SimpleDQN/'
import GreedyClient

class DQNClient(WebSocketClient):

    def __init__(self, url, id, agent=None, agent_list=None, is_eval=False, load_check_point_name=None):
        super().__init__(url)
        self.agent_id = id
        self.state = State()
        self.action = bucket_rule.Action()
        self.is_eval = is_eval
        self.cnt = 0
        self.eval_reward=0
        if load_check_point_name is None:
            self.check_point_name = save_dir+'agent'+str(id)
        else:
            self.check_point_name = load_check_point_name
        # Rank info
        self.current_rank = None
        self.self_rank = None
        self.opponent_rank = None

        # position info
        self.first_act_player = None
        self.position = None

        # hand info
        self.initial_hand = None
        self.action_history = [[],[],[],[]]
        self.card_played = np.zeros((4,16))

        # reward related
        # indicate if self is the first of the round
        self.self_start = False

        # agent
        # self.sess = tf.Session()
        if agent is None:
            self.agent = dqn.DQN(
                sess,
                player_id=self.agent_id,
                state_representation_size=108,
                num_actions=118,
                learning_rate=1e-3,
                replay_buffer_capacity=10000,
                epsilon_decay_duration=2000,
                hidden_layers_sizes=[128,256,256,128],
                update_target_network_every=250,
                batch_size=64
            )

            if os.path.exists(self.check_point_name):
                self.agent.restore(self.check_point_name)
                print('load successfully\n')
            else:
                if not os.path.exists(save_dir):
                    os.mkdir(save_dir)
                os.mkdir(self.check_point_name)
        else:
            self.agent = agent
        self.agent_list = agent_list


    def new_start(self):
        self.current_rank = None
        self.self_rank = None
        self.opponent_rank = None
        self.first_act_player = None
        self.action_history = [[],[],[],[]]
        self.initial_hand = None
        self.position = None
        self.card_played = np.zeros((4,16))

    def set_eval(self):
        self.is_training = False

    def set_train(self):
        self.is_training = True

    def opened(self):
        pass

    def closed(self, code, reason=None):
        print("Closed down", code, reason)

    def received_message(self, message):
        message = json.loads(str(message))                                    # 先序列化收到的消息，转为Python中的字典
        result = self.state.parse(message)                                             # 调用状态对象来解析状态
        stage = result[0]
        m_type = result[1]
        if m_type == 'notify':
            if stage == 'beginning':
                self.new_start()
                self.initial_hand = copy.deepcopy(result[2])
                assert type(self.initial_hand) == list
                self.position = result[3]
                assert type(self.position) == int
                if self.agent_id != self.position:
                    assert self.is_eval
                    self.agent_id = self.position
                    self.agent = self.agent_list[self.position]
            elif stage == 'play':
                pos = result[2]
                if self.first_act_player is None:
                    self.first_act_player = pos
                action_card = result[3]
                action_copy = copy.deepcopy(action_card)
                if type(action_copy) == str:
                    convert = action_copy.replace("'", '"')
                    action_copy = json.loads(convert)[-1]
                else:
                    action_copy = action_copy[-1]
                    # action_copy = my_utils.convert_his_action2normal(action_copy)[2]
                self.action_history[pos].append(action_copy)
                # self.record_action(action_copy, pos)
            elif stage == 'tribute':
                pass
            elif stage == 'anti-tribute':
                pass
            elif stage == 'back':
                pass
            elif stage == 'episodeOver':
                winner_order = result[2]
                rest_card = result[4]
                if self.position != winner_order[-1]:
                    our_cards = []
                else:
                    our_cards = []
                    for tmp in rest_card:
                        if tmp[0] == self.position:
                            our_cards = tmp[-1]
                observe_info = self.get_observed_info()
                if len(our_cards) == 0:
                    hand_info = np.zeros(27)
                else:
                    hand_info = my_utils.cards2array(our_cards, self.current_rank)
                info = np.concatenate([hand_info,observe_info])

                time_step = self.get_time_step(info, None, order=winner_order,end=True )
                if self.is_eval:
                    self.eval_reward += winner_order.index(self.position) + winner_order.index((self.position+2)%4) + 2
                    print(self.cnt,',my_pos: ',self.position ,winner_order)
                self.agent.step(time_step)
                self.cnt += 1
                if self.cnt % 10 == 0 and not self.is_eval:
                    print(time.strftime('%Y.%m.%d %H:%M:%S',time.localtime()),' : ',self.cnt)
                if self.cnt % 128 == 0 and not self.is_eval:
                    self.agent.save(self.check_point_name)
                    print('Save,',self.cnt)
                # TODO: add last time step to buffer
            elif stage == 'gameOver':
                cur_times = result[2]
                if cur_times %10 == 0 and not self.is_eval:
                    self.agent.save(self.check_point_name)
                pass
            elif stage == 'gameResult':
                if not self.is_eval:
                    self.agent.save(self.check_point_name)
                else:
                    print(self.get_eval_reward()/ self.cnt)
                pass
            else:
                raise Exception("unknown stage:", stage)


        else:
            assert m_type == 'act'
            if stage == 'play':
                current_hand = result[2]
                # print(current_hand)

                legal_action_list = result[3]
                legal_action_index_list = []
                if self.self_rank is None:
                    assert self.opponent_rank is None and self.current_rank is None
                    self.self_rank = my_utils.rank2int(result[4])
                    self.opponent_rank = my_utils.rank2int(result[5])
                    self.current_rank = my_utils.rank2int(result[6])
                else:
                    assert self.self_rank == my_utils.rank2int(result[4])
                    assert self.opponent_rank == my_utils.rank2int(result[5])
                    assert self.current_rank == my_utils.rank2int(result[6])
                greater_pos = result[7]
                greater_action = result[8][-1]

                for player_id in range(4):
                    if len(self.action_history[player_id]) != 0:
                        # [-1] is the cards in last action
                        self.record_action(self.action_history[player_id][-1], player_id)
                # 27
                hand_info = my_utils.cards2array(current_hand, self.current_rank)
                # 81
                observe_info = self.get_observed_info(greater_pos, greater_action)
                for card_action in legal_action_list:
                    action_type = card_action[0]
                    action_rank = card_action[1]
                    action_card = card_action[2]
                    '''if action_rank == 'B' or action_rank == 'R':
                        temp = ['Trips','ThreeWithTwo', 'TwoTrips']
                        if action_type in temp:
                            continue'''
                    if action_type == 'Straight' or action_type == 'StraightFlush':
                        if action_rank == 'J' or action_rank == 'Q' or action_rank == 'K' or action_rank == '':

                            continue
                    action_index = my_utils.action2index(action_type, action_rank, action_card, self.current_rank)
                    if action_index not in legal_action_index_list:
                        legal_action_index_list.append(action_index)
                info = np.concatenate([hand_info,observe_info])
                if len(self.action_history[self.position]) ==0:
                    first = True
                else:
                    first = False
                time_step = self.get_time_step(info,legal_action_index_list,start=first)
                agent_output = self.agent.step(time_step,is_evaluation=self.is_eval)
                '''if np.random.choice(2 ,p=[0.999, 0.001]) and not self.is_eval:
                    t = np.zeros(118)
                    t[legal_action_index_list] = 1
                    print('action_prob:',agent_output.probs)
                    print('legal:',t)
                    print('action:', legal_action_list,'\n')'''
                action_bucket_index = agent_output.action
                if self.is_eval:
                    probs = agent_output.probs
                    action_bucket_index = int(probs.argmax())
                #return_index = my_utils.random_choose_action(action_bucket_index, self.current_rank, legal_action_list)
                # return_index = greedy_rule.get_index(action_bucket_index,greater_pos,self.current_rank,legal_action_list)
                bucket_choose = bucket_rule.Action()
                return_index = bucket_choose.parse(message, bucket_id=action_bucket_index)
                # print('choose action',return_index)
                self.send(json.dumps({"actIndex": int(return_index)}))
                pass
            elif stage == 'tribute':
                act_index = self.action.parse(message)
                self.send(json.dumps({"actIndex": act_index}))
            elif stage == 'back':
                act_index = self.action.parse(message)
                self.send(json.dumps({"actIndex": act_index}))
            else:
                raise Exception('Unknown stage: ', stage)


        # if "actionList" in message:                                           # 需要做出动作选择时调用动作对象进行解析
        #    act_index = self.action.parse(message)
        #    self.send(json.dumps({"actIndex": act_index}))
    def get_time_step(self, current_info, legal_action, order=None, start=False, end=False):
        observations = {
            "info_state": [],
            "legal_actions": [],
            "current_player": [],
            "serialized_state": []
        }
        rewards = [0, 0, 0, 0]
        if not start:
            if self.action_history[self.position][-1] == 'PASS':
                rewards[self.position] = -0.01


        for id in range(4):
            observations["info_state"].append(None if id != self.position else current_info)
            observations["legal_actions"].append(None if id != self.position else legal_action)
        observations["current_player"] = self.position
        if order is not None:
            rewards = my_utils.get_reward(order)

        discounts = [1.0, 1.0, 1.0, 1.0]
        if start:
            step_type = rl_environment.StepType.FIRST
        elif end:
            step_type = rl_environment.StepType.LAST
            discounts = [0.0] * 4
        else:
            step_type = rl_environment.StepType.MID

        return rl_environment.TimeStep(
            observations=observations,
            rewards=rewards,
            discounts=discounts,
            step_type=step_type)

    def record_action(self, action_card, player_id):
        if action_card == 'PASS':
            return
        card_array = my_utils.cards2array(action_card, self.current_rank, only_card=True)
        self.card_played[player_id] += card_array

    def get_observed_info(self, greater_pos=None, greater_action=None):
        """

        Return:
             81-
        """
        # flatten = np.reshape(self.card_played, -1)
        cards_played = np.zeros([4, 16])
        for id in range(4):
            player_id = (self.position + id) % 4
            cards_played[id] = self.card_played[player_id]

            '''if len(self.action_history[id]) == 0:
                pass
            else:
                card_this_round[player_id]=my_utils.cards2array(self.action_history[id][-1], self.current_rank, only_card=True)'''
        if greater_pos is None:
            return np.concatenate([cards_played.flatten(),np.zeros(17)])
        greater_pos_array = np.zeros(1)
        if type(greater_pos)==str:
            greater_pos = int(greater_pos)
        temp = (greater_pos - self.position) % 4
        greater_pos_array[0] = temp
        greater_action_array = my_utils.cards2array(greater_action, self.current_rank, only_card=True) if greater_action is not None else np.zeros(16)

        res = np.concatenate([cards_played.flatten(), greater_pos_array,greater_action_array])
        #assert False, str(len(res))+','+str(len(flatten))+','+str(len(card_this_round.flatten()))
        return res

    def reset_eval_reward(self):
        self.eval_reward = 0

    def get_eval_reward(self):
        return self.eval_reward


class MyThread(threading.Thread):
    def __init__(self, url, agent_id, agent=None, agent_lists=None, is_eval=False):
        threading.Thread.__init__(self)
        self.threadID = agent_id
        self.url = url
        if not is_eval:
            try:
                self.client = DQNClient(url,agent_id, agent=agent)
            except Exception:
                print("#############",self.threadID)
                raise Exception
        else:
            if self.threadID == 1 or self.threadID == 3:
                self.client = opponent.ExampleClient(url)
                # self.client = rdClient.RandomClient(url)
                # self.client = AI.ExampleClient(url)
            else:
                #assert agent is not None, 'should specify agent when test!\n'
                self.client = DQNClient(url, agent_id, agent=agent, agent_list=agent_lists, is_eval=True)
        # self.file = open(str(self.threadID)+'.txt','w')

    def run(self):
        print('start %d\n' % self.threadID)
        try:
            if use_greedy:
                            threadLock.acquire()

                            if len(greedy_agent_list) == 0:
                                new_sess = tf.Session()

                                for i in range(4):
                                    greedy_agent_list.append(GreedyClient.load_agent(agent_names[i], i, session=new_sess))
                            threadLock.release()
                            print(len(greedy_agent_list))
                            print(greedy_agent_list)

                            if self.threadID == 1 or self.threadID == 3:
                                self.client = GreedyClient.GreedyClient(self.url,self.threadID,agent=greedy_agent_list[self.threadID], agent_list=greedy_agent_list, is_eval=True)
                            print(self.threadID,self.client.agent)
            self.client.connect()
            self.client.run_forever()
        except Exception:
            self.client.close()


def load_agent(agent_name, agent_id, session=None):
    if session is None:
        session = tf.Session
    agent = dqn.DQN(
        sess,
        player_id=agent_id,
        state_representation_size=108,
        num_actions=118,
        learning_rate=1e-3,
        replay_buffer_capacity=10000,
        epsilon_decay_duration=2000,
        hidden_layers_sizes=[128, 256, 256, 128],
        update_target_network_every=250,
        batch_size=64
    )
    load_dir = save_dir+agent_name
    print('loading', load_dir)

    agent.restore(load_dir)
    print('load ', load_dir)

    return agent


if __name__ == '__main__':
    Test = True
    use_greedy = True
    threadLock = threading.Lock()
    greedy_agent_list = []
    agent_names = ['agent0', 'agent1', 'agent2', 'agent3']

    urls = ['ws://127.0.0.1:23456/game/client0', 'ws://127.0.0.1:23456/game/client1',
            'ws://127.0.0.1:23456/game/client2', 'ws://127.0.0.1:23456/game/client3']
    with tf.Session() as sess:
        if Test:
            agent_names = ['agent0', 'agent1', 'agent2', 'agent3']
            agent_list = []
            for i in range(4):
                agent_list.append(load_agent(agent_names[i],i,session=sess))
            thread = []
            for i in range(4):
                new_thread = MyThread(urls[i], i,agent=agent_list[i], agent_lists=agent_list, is_eval=True)
                thread.append(new_thread)
                new_thread.start()
                time.sleep(1)
            for i in thread:
                i.join()
                time.sleep(1)
            print('end')

            pass
        else:
            thread = []
            for i in range(4):
                new_thread = MyThread(urls[i], i)
                thread.append(new_thread)
                new_thread.start()
            for i in thread:
                i.join()
            print('end')
