from ..Agent import Agent
import numpy as np
from .KuhnInfo import Cards, KUHN_TREE, KuhnAction

class KuhnAgent(Agent):
    card: Cards

    def _compute_equilibrium(self, private: bool):
        own_action = self._start
        transitions = KUHN_TREE.transitions
        for action in self._state.action_history:
            if own_action and private:
                transitions = [t for t in transitions if t.card == self.card]
            transitions = [t for t in transitions if t.action == action]
            new_transitions = []
            for t in transitions:
                new_transitions += t.node.transitions
            transitions = new_transitions
            own_action = not own_action
        if private:
            transitions = [t for t in transitions if t.card == self.card]
        result = np.zeros((self._state.actions_number))
        for action in list(KuhnAction):
            action_equilibrium = [
                t.equilibrium for t in transitions if t.action == action]
            if len(action_equilibrium) > 0:
                result[action.value] = sum(
                    action_equilibrium) / len(action_equilibrium)
            else:
                result[action.value] = 0
        return result

    def compute_private_equilibrium(self):
        return self._compute_equilibrium(True)

    def compute_public_equilibrium(self):
        return self._compute_equilibrium(False)

    def has_bet(self):
        return KuhnAction.BET in self.actions

    def has_fold(self):
        return KuhnAction.FOLD in self.actions
