import math
import torch
import numpy as np

class Node:

    def __init__(self, game, args, state, parent = None, action = None, prob = 0, visits = 0):
        self.game = game
        self.args = args
        self.state = state
        self.parent = parent
        self.action = action
        self.prob = prob

        self.children = []
        self.visits = visits
        self.value = 0

    def leaf_or_not(self):
        return len(self.children) > 0

    def search(self):
        best_child = None
        best_ucb = -np.inf
        for child in self.children:
            ucb = self.get_ucb(child)
            if best_ucb < ucb:
                best_ucb = ucb
                best_child = child
        return best_child

    def get_ucb(self, child):
        if child.visits == 0:
            q_value = 0
        else:
            q_value = 1 - ((child.value / child.visits) + 1) / 2

        return q_value + self.args['EXPLORATION_CONSTANT'] * (math.sqrt(self.visits) / (child.visits + 1)) * child.prob

    def expand(self, policy):

        for move, prob in enumerate(policy):
            if prob > 0:
                child = self.state.copy()
                child = self.game.make_move(child, move, 1)
                if self.args["ADVERSARIAL"]:
                    child = self.game.change_perspective(child, player = -1)

                child = Node(self.game, self.args, child, self, move, prob)
                self.children.append(child)

    def backpropagate(self,state_value):
        self.value += state_value
        self.visits += 1
        if self.args["ADVERSARIAL"]:
            state_value = self.game.get_opponent_value(state_value)
        if self.parent is not None:
            self.parent.backpropagate(state_value)


class Alpha_MCTS:

    def __init__(self, game, args, model):
        self.game = game
        self.args = args
        self.model = model

    @torch.no_grad()
    def search(self, states, spGames):

        policy, _ = self.model(
                        torch.tensor(self.game.get_encoded_state(states), device = self.model.device
                    ))

        policy = torch.softmax(policy, axis = 1).cpu().numpy()

        if self.args["ROOT_RANDOMNESS"]:
            policy = (1 - self.args["DIRICHLET_EPSILON"]) * policy + self.args["DIRICHLET_EPSILON"] * np.random.dirichlet([self.args["DIRICHLET_ALPHA"]] * self.game.possible_state, size = policy.shape[0])

        for i, spg in enumerate(spGames):
            spg_policy = policy[i]
            valid_state = self.game.get_valid_moves(states[i])
            spg_policy *= valid_state
            spg_policy /= np.sum(spg_policy)

            spg.root = Node(self.game, self.args, states[i], visits = 1)
            spg.root.expand(spg_policy)

        for _ in range(self.args["NO_OF_SEARCHES"]):

            for spg in spGames:
                spg.node = None
                node = spg.root

                while node.leaf_or_not():
                    node = node.search()

                is_terminal, value = self.game.know_terminal_value(node.state, node.action)

                if self.args["ADVERSARIAL"]:
                    value = self.game.get_opponent_value(value)

                if is_terminal:
                    node.backpropagate(value)
                else:
                    spg.node = node

            expandabel_spgs = [mapping_index for mapping_index in range(len(spGames)) if spGames[mapping_index].node is not None]

            if len(expandabel_spgs) > 0:
                states = np.stack([spGames[mapping_index].node.state for mapping_index in expandabel_spgs])
                policy, value = self.model(
                    torch.tensor(self.game.get_encoded_state(states), device = self.model.device)
                )

                policy = torch.softmax(policy, axis = 1).cpu().numpy().astype(np.float64)
                value = value.cpu().numpy()

            for i, mapping_index in enumerate(expandabel_spgs):
                node = spGames[mapping_index].node
                spg_policy, spg_value = policy[i], value[i]
                valid_state = self.game.get_valid_moves(node.state)

                spg_policy *= valid_state
                spg_policy /= np.sum(spg_policy)

                node.expand(spg_policy)
                node.backpropagate(spg_value)