import numpy as np
from collections import defaultdict
from abc import ABC, abstractmethod
class RealtimeSearchMCTSNode(ABC):

    def __init__(self, state, parent=None):
        '''
        Parameters
        ----------
        state : poker_env.state.TwoPlayersAbstractGameState
            require:
                __init__
                game_result(get_payoffs)
                is_over(OK)
                step(OK)
                get_legal_actions(OK)
        parent : RealtimeSearchMCTSNode
        '''
        self.state = state
        self.parent = parent
        self.children = []

    @property
    @abstractmethod
    def untried_actions(self):
        '''

        Returns:
            (list): list of poker_env.AbstractAction

        '''
        pass

    @property
    @abstractmethod
    def q(self):
        pass

    @property
    @abstractmethod
    def n(self):
        pass

    @abstractmethod
    def expand(self):
        pass

    @abstractmethod
    def is_terminal_node(self):
        pass

    @abstractmethod
    def rollout(self):
        pass

    @abstractmethod
    def backpropagate(self, reward):
        pass

    def is_fully_expanded(self):
        return len(self.untried_actions) == 0

    def best_child(self, c_param=1.4):
        choices_weights = [
            (c.q / c.n) + c_param * np.sqrt((2 * np.log(self.n) / c.n))
            for c in self.children
        ]
        return self.children[np.argmax(choices_weights)]

    def rollout_policy(self, possible_actions):
        '''
        TODO:
        random by  probability distribution
        '''        
        return possible_actions[np.random.randint(len(possible_actions))]

class NoLimitTexasHoldemGameMCTSNode(RealtimeSearchMCTSNode):

    def __init__(self, state, parent=None):
        super().__init__(state, parent)
        self._number_of_visits = 0.
        self._results = defaultdict(int)  
        '''
        TODO:adjust the structure
        '''
        self._untried_actions = None

    @property
    def untried_actions(self):
        if self._untried_actions is None:
            self._untried_actions = self.state.get_legal_actions()
        return self._untried_actions

    @property
    def q(self):
        '''
        TODO:adjust the compute
        '''
        wins = self._results[self.parent.state.next_to_move]
        loses = self._results[-1 * self.parent.state.next_to_move]
        return wins - loses

    @property
    def n(self):
        return self._number_of_visits

    def expand(self):
        action = self.untried_actions.pop()
        next_state, _ = self.state.step(action)
        child_node = NoLimitTexasHoldemGameMCTSNode(
            next_state, parent=self
        )
        self.children.append(child_node)
        return child_node

    def is_terminal_node(self):
        return self.state.is_over()

    def rollout(self):
        current_rollout_state = self.state
        while not current_rollout_state.is_over():
            possible_actions = current_rollout_state.get_legal_actions()
            action = self.rollout_policy(possible_actions)
            current_rollout_state = current_rollout_state.step(action)
        return current_rollout_state.game_result

    def backpropagate(self, result):
        self._number_of_visits += 1.
        self._results[result] += 1.
        '''
        '''
        if self.parent:
            self.parent.backpropagate(result)