import json
import random
import numpy as np
import my_utils as utils
from ws4py.client.threadedclient import WebSocketClient
from open_spiel.python import rl_environment
from open_spiel.python.algorithms import policy_gradient
import tensorflow.compat.v1 as tf
import threading

class AIClient(WebSocketClient):

    def __init__(self, url, agent_id=0, info_state_size=0, num_actions=0 ):
        super().__init__(url)
        self.id = agent_id
        self.sess = tf.Session()
        self.agent = policy_gradient.PolicyGradient(
            self.sess,
            player_id=agent_id,
            info_state_size=info_state_size,
            num_actions=num_actions,
            loss_str="rpg"

        )
        self.initial_hand = None
        self.current_hand = None
        self.game_init = None
        self.action_list = None

    def opened(self):
        pass
        # self.send(json.dumps({"name": self.name}))w

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

    def received_message(self, message):
        content = json.loads(str(message))
        print(content)
        if "action_list" in content and content["action_list"]:
            card_type = random.choice(list(content["action_list"].keys()))
            print("Choose type ", card_type)
            rank = random.choice(list(content["action_list"][card_type].keys()))
            action = random.choice(list(content["action_list"][card_type][rank]))
            print("Choose action:", action)
            self.send(json.dumps({"action": action, "type": card_type, "rank": rank}))

    def new_game(self):
        self.initial_hand = None
        self.current_hand = None
        self.game_init = True
        self.action_list = None

    def extract_action(self, content):
        pass

    def get_init_hand(self, content):
        assert self.game_init == True
        hand = content['hand_cards']
        return hand

    def get_time_step(self, content, initial_hand, action_list):
        """

        Returns:
            A `TimeStep` namedtuple containing:
                observation: list of dicts containing one observations per player, each
                    corresponding to `observation_spec()`.
                reward: list of rewards at this timestep, or None if step_type is
                    `StepType.FIRST`.
                discount: list of discounts in the range [0, 1], or None if step_type is
                    `StepType.FIRST`.
                step_type: A `StepType` value.
        """
        is_abstract = False
        observations = {
            "info_state": [],
            "legal_actions": [],
            "current_player": [],
            "serialized_state": []
        }
        rewards = []
        if content['type'] == 0:
            # end
            step_type = rl_environment.StepType.LAST
            self._end_game = True
        elif content['type'] == 2:
            # need to act
            legal_actions = content['action_list']
            self_rank = content['self_rank']
            opponent_rank = content['opponent_rank']
            current_rank = content['current_rank']
            last_player_action = content['action_performed']
            new_action_list = utils.abstract(legal_actions,is_abstract)

            pass
        else:
            raise Exception("No need to act\n")

    @staticmethod
    def get_current_player_id(content):
        #Todo: need to make sure the 'current player' in the content is the
        id = content['current_player']
        return (id+3)%4

    @staticmethod
    def get_reward(content):
        rewards = np.zeros(4)
        if 'winners' in content:
            if content['winners'] is None:
                raise Exception
            winners = content['winners']
            winner = winners[0]

            assert type(winner) == int
            if winner == 1 or winner == 3:
                rank = 4 - winners.index(1) - winners.index(3)
                r = [rank, rank, -rank, -rank]
                pos = [1,3,0,2]
            else:
                rank = 4 - winners.index(0) - winners.index(2)
                r = [rank, rank, -rank, -rank]
                pos = [0,2,1,3]
            rewards[pos] = r
        return list(rewards)


class MyThread(threading.Thread):
    def __init__(self, url):
        threading.Thread.__init__(self)
        self.threadID = int(url[-1])
        self.client = AIClient(url,self.threadID)

    def run(self):
        print('start %d\n' % self.threadID)
        try:
            self.client.connect()
            self.client.run_forever()
        except Exception:
            self.client.close()





if __name__ == '__main__':
    urls = 'ws://127.0.0.1:23456/game/client'
    thread = []
    for i in range(4):
        new_thread = MyThread(urls+str(i))
        thread.append(new_thread)
        new_thread.start()
    for i in thread:
        i.join()
    print('end')

