import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch_geometric.data import Data
from torch_geometric.nn import GATConv
from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T

name_data = 'Cora'
dataset = Planetoid(root='./data/', name=name_data)

from torch_geometric.utils import to_networkx
cora = to_networkx(dataset.data)

print(cora.is_directed())

node_classes = dataset.data.y.data.numpy()
print(node_classes)

node_color = ["red","blue","green","yellow","peru","violet","cyan"] 
node_label = np.array(list(cora.nodes)) 

import matplotlib.pyplot as plt
import networkx as nx

pos = nx.layout.spring_layout(cora)

plt.figure(figsize=(16,12))

for i in np.arange(len(np.unique(node_classes))):
    node_list = node_label[node_classes == i]
    nx.draw_networkx_nodes(cora, pos, nodelist=list(node_list),
                           node_size=50,  
                           node_color=node_color[i], 
                           alpha=0.8)

nx.draw_networkx_edges(cora, pos,width=1,edge_color="black")
plt.show()

import hvplot.networkx as hvnx
options = {
    'width': 800,
    'height': 1000
}

plt_nodes = []
for i in np.arange(len(np.unique(node_classes))):
    nodelist = node_label[node_classes == i]
    plt = hvnx.draw_networkx_nodes(cora, pos, nodelist=list(nodelist), node_color=node_color[i], **options)
    plt_nodes.append(plt)

plt_edges = hvnx.draw_networkx_edges(cora, pos, arrowstyle='->', edge_width=2, colorbar=True, **options)

import functools
import operator
plt_edges * functools.reduce(operator.mul, plt_nodes)

class TreeNode:
    """
    MCTS Tree Node
    """

    c_puct: ClassVar[int] = 5  # class-wise global param c_puct, exploration weight factor.

    _parent: TreeNode
    _children: Dict[int, TreeNode]  # map from action to TreeNode
    _visit_num: int
    _Q: float   # Q value of the node, which is the mean action value.
    _prior: float

class TreeNode:

    def get_puct(self) -> float:
        """
        Computes AlphaGo Zero PUCT (polynomial upper confidence trees) of the node.

        :return: Node PUCT value.
        """
        U = (TreeNode.c_puct * self._prior * np.sqrt(self._parent._visit_num) / (1 + self._visit_num))
        return self._Q + U

class TreeNode:

    def select(self) -> Tuple[Pos, TreeNode]:
        """
        Selects an action(Pos) having max UCB value.

        :return: Action and corresponding node
        """
        return max(self._children.items(), key=lambda act_node: act_node[1].get_puct())

class TreeNode:

    def propagate_to_root(self, leaf_value: float):
        """
        Updates current node with observed leaf_value and propagates to root node.

        :param leaf_value:
        :return:
        """
        if self._parent:
            self._parent.propagate_to_root(-leaf_value)
        self._update(leaf_value)

    def _update(self, leaf_value: float):
        """
        Updates the node by newly observed leaf_value.

        :param leaf_value:
        :return:
        """
        self._visit_num += 1
        # new Q is updated towards deviation from existing Q
        self._Q += 0.5 * (leaf_value - self._Q)

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def _next_step_play_act_probs(self, game: ConnectNGame) -> Tuple[List[Pos], ActionProbs]:
        """
        For the given game status, run playouts number of times specified by self._playout_num.
        Returns the action distribution according to AlphaGo Zero MCTS play formula.

        :param game:
        :return: actions and their probability
        """

        for n in range(self._playout_num):
            self._playout(copy.deepcopy(game))

        act_visits = [(act, node._visit_num) for act, node in self._current_root._children.items()]
        acts, visits = zip(*act_visits)
        act_probs = softmax(1.0 / MCTSAlphaGoZeroPlayer.temperature * np.log(np.array(visits) + 1e-10))

        return acts, act_probs

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def get_action(self, board: PyGameBoard) -> Pos:
        """
        Method defined in BaseAgent.

        :param board:
        :return: next move for the given game board.
        """
        return self._get_action(copy.deepcopy(board.connect_n_game))[0]

    def _get_action(self, game: ConnectNGame) -> Tuple[MoveWithProb]:
        epsilon = 0.25
        avail_pos = game.get_avail_pos()
        move_probs: ActionProbs = np.zeros(game.board_size * game.board_size)
        assert len(avail_pos) > 0

        # the pi defined in AlphaGo Zero paper
        acts, act_probs = self._next_step_play_act_probs(game)
        move_probs[list(acts)] = act_probs
        if self._is_training:
            # add Dirichlet Noise when training in favour of exploration
            p_ = (1-epsilon) * act_probs + epsilon * np.random.dirichlet(0.3 * np.ones(len(act_probs)))
            move = np.random.choice(acts, p=p_)
            assert move in game.get_avail_pos()
        else:
            move = np.random.choice(acts, p=act_probs)

        self.reset()
        return move, move_probs

class MCTSAlphaGoZeroPlayer(BaseAgent):

    def self_play_one_game(self, game: ConnectNGame) \
            -> List[Tuple[NetGameState, ActionProbs, NDArray[(Any), np.float]]]:
        """

        :param game:
        :return:
            Sequence of (s, pi, z) of a complete game play. The number of list is the game play length.
        """

        states: List[NetGameState] = []
        probs: List[ActionProbs] = []
        current_players: List[np.float] = []

        while not game.game_over:
            move, move_probs = self._get_action(game)
            states.append(convert_game_state(game))
            probs.append(move_probs)
            current_players.append(game.current_player)
            game.move(move)

        current_player_z = np.zeros(len(current_players))
        current_player_z[np.array(current_players) == game.game_result] = 1.0
        current_player_z[np.array(current_players) == -game.game_result] = -1.0
        self.reset()

        return list(zip(states, probs, current_player_z))

class MCTSAlphaGoZeroPlayer(BaseAgent):
    
    def _playout(self, game: ConnectNGame):
        """
        From current game status, run a sequence down to a leaf node, either because game ends or unexplored node.
        Get the leaf value of the leaf node, either the actual reward of game or action value returned by policy net.
        And propagate upwards to root node.

        :param game:
        """
        player_id = game.current_player

        node = self._current_root
        while True:
            if node.is_leaf():
                break
            act, node = node.select()
            game.move(act)

        # now game state is a leaf node in the tree, either a terminal node or an unexplored node
        act_and_probs: Iterator[MoveWithProb]
        act_and_probs, leaf_value = self._policy_value_net.policy_value_fn(game)

        if not game.game_over:
            # case where encountering an unexplored leaf node, update leaf_value estimated by policy net to root
            for act, prob in act_and_probs:
                game.move(act)
                child_node = node.expand(act, prob)
                game.undo()
        else:
            # case where game ends, update actual leaf_value to root
            if game.game_result == ConnectNGame.RESULT_TIE:
                leaf_value = ConnectNGame.RESULT_TIE
            else:
                leaf_value = 1 if game.game_result == player_id else -1
            leaf_value = float(leaf_value)

        # Update leaf_value and propagate up to root node
        node.propagate_to_root(-leaf_value)

NetGameState = NDArray[(4, Any, Any), np.int]


def convert_game_state(game: ConnectNGame) -> NetGameState:
    """
    Converts game state to type NetGameState as ndarray.

    :param game:
    :return:
        Of shape 4 * board_size * board_size.
        [0] is current player positions.
        [1] is opponent positions.
        [2] is last move location.
        [3] all 1 meaning move by black player, all 0 meaning move by white.
    """
    state_matrix = np.zeros((4, game.board_size, game.board_size))

    if game.action_stack:
        actions = np.array(game.action_stack)
        move_curr = actions[::2]
        move_oppo = actions[1::2]
        for move in move_curr:
            state_matrix[0][move] = 1.0
        for move in move_oppo:
            state_matrix[1][move] = 1.0
        # indicate the last move location
        state_matrix[2][actions[-1]] = 1.0
    if len(game.action_stack) % 2 == 0:
        state_matrix[3][:, :] = 1.0  # indicate the colour to play
    return state_matrix[:, ::-1, :]

def backward_step(self, state_batch: List[NetGameState], probs_batch: List[ActionProbs],
               value_batch: List[NDArray[(Any), np.float]], lr) -> Tuple[float, float]:
    if self.use_gpu:
        state_batch = Variable(torch.FloatTensor(state_batch).cuda())
        probs_batch = Variable(torch.FloatTensor(probs_batch).cuda())
        value_batch = Variable(torch.FloatTensor(value_batch).cuda())
    else:
        state_batch = Variable(torch.FloatTensor(state_batch))
        probs_batch = Variable(torch.FloatTensor(probs_batch))
        value_batch = Variable(torch.FloatTensor(value_batch))

    self.optimizer.zero_grad()
    for param_group in self.optimizer.param_groups:
        param_group['lr'] = lr

    log_act_probs, value = self.policy_value_net(state_batch)
    # loss = (z - v)^2 - pi*T * log(p) + c||theta||^2
    value_loss = F.mse_loss(value.view(-1), value_batch)
    policy_loss = -torch.mean(torch.sum(probs_batch * log_act_probs, 1))
    loss = value_loss + policy_loss
    loss.backward()
    self.optimizer.step()
    entropy = -torch.mean(torch.sum(torch.exp(log_act_probs) * log_act_probs, 1))
    return loss.item(), entropy.item()