"""Policy gradient learning."""
import numpy as np
import torch
import torch.nn.functional as F
from torch.optim import SGD

from tugo.agents.base import Agent
from tugo.agents.helpers import is_point_an_eye
from game_logic import goboard


def policy_gradient_loss(y_true, y_pred):
    clip_pred = torch.clamp(y_pred, 1e-10, 1 - 1e-10)
    loss = -1 * y_true * torch.log(clip_pred)
    return torch.mean(torch.sum(loss, dim=1))


class PolicyAgent(Agent):
    """An agent that uses a deep policy network to select moves."""
    def __init__(self, model, encoder):
        super().__init__()
        self._model = model
        self._encoder = encoder
        self._collector = None
        self._temperature = 0.0

    def predict(self, game_state):
        encoded_state = self._encoder.encode(game_state)
        input_tensor = torch.tensor([encoded_state], dtype=torch.float32).to('cuda')
        with torch.no_grad():
            output_tensor = self._model(input_tensor)
        return output_tensor.cpu().numpy()[0]

    def set_temperature(self, temperature):
        self._temperature = temperature

    def set_collector(self, collector):
        self._collector = collector

    def select_move(self, game_state):
        num_moves = self._encoder.board_width * self._encoder.board_height

        move_probs = self.predict(game_state)

        move_probs = move_probs ** 3
        eps = 1e-6
        move_probs = np.clip(move_probs, eps, 1 - eps)
        move_probs = move_probs / np.sum(move_probs)

        candidates = np.arange(num_moves)
        ranked_moves = np.random.choice(candidates, num_moves, replace=False, p=move_probs)

        for point_idx in ranked_moves:
            point = self._encoder.decode_point_index(point_idx)
            if game_state.is_valid_move(goboard.Move.play(point)) and \
                    not is_point_an_eye(game_state.board, point, game_state.next_player):
                if self._collector is not None:
                    self._collector.record_decision(state=board_tensor, action=point_idx)
                return goboard.Move.play(point)

        return goboard.Move.pass_turn()

    def train(self, experience, lr=1e-7, clipnorm=1.0, batch_size=512):
        opt = SGD(self._model.parameters(), lr=lr)
        n = experience.states.shape[0]
        num_moves = self._encoder.board_width * self._encoder.board_height
        y = torch.zeros((n, num_moves))
        for i in range(n):
            action = experience.actions[i]
            reward = experience.rewards[i]
            y[i][action] = reward

        for epoch in range(1):
            permutation = torch.randperm(n)
            for i in range(0, n, batch_size):
                indices = permutation[i:i+batch_size]
                batch_x, batch_y = experience.states[indices], y[indices]
                opt.zero_grad()
                outputs = self._model(batch_x)
                loss = F.cross_entropy(outputs, batch_y)
                loss.backward()
                opt.step()

    def save(self, file_path):
        torch.save(self._model.state_dict(), file_path)

    @classmethod
    def load(cls, file_path, encoder):
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = AlphaGoModel(encoder.get_input_shape(), is_policy_net=True).to(device)
        model.load_state_dict(torch.load(file_path))
        return cls(model, encoder)