import torch
import torch.nn.functional as F
from torch import optim

from tugo import encoders
from tugo.agents.base import Agent
from tugo.agents.helpers import is_point_an_eye
# from tugo.game_logic import goboard
from tugo.game_logic import goboard_fast as goboard


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


def normalize(x):
    total = torch.sum(x)
    return x / total


class PolicyAgent(Agent):
    def __init__(self, model, encoder):
        Agent.__init__(self)
        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])
        return self._model(input_tensor)[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

        board_tensor = self._encoder.encode(game_state)
        x = torch.tensor([board_tensor])

        if torch.rand(1).item() < self._temperature:
            move_probs = torch.ones(num_moves) / num_moves
        else:
            move_probs = self._model(x)[0]

        eps = 1e-5
        move_probs = torch.clamp(move_probs, eps, 1 - eps)
        move_probs = move_probs / torch.sum(move_probs)

        candidates = torch.arange(num_moves)
        ranked_moves = torch.multinomial(move_probs, num_moves, replacement=False)
        for point_idx in ranked_moves:
            point = self._encoder.decode_point_index(point_idx.item())
            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.item()
                    )
                return goboard.Move.play(point)
        return goboard.Move.pass_turn()

    def train(self, experience, lr=0.0000001, clipnorm=1.0, batch_size=512):
        opt = optim.SGD(self._model.parameters(), lr=lr, clipnorm=clipnorm)

        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

        self._model.train()
        opt.zero_grad()
        outputs = self._model(experience.states)
        loss = F.cross_entropy(outputs, y)
        loss.backward()
        opt.step()

    def serialize(self, h5file):
        h5file.create_group('encoder')
        h5file['encoder'].attrs['name'] = self._encoder.name()
        h5file['encoder'].attrs['board_width'] = self._encoder.board_width
        h5file['encoder'].attrs['board_height'] = self._encoder.board_height
        h5file.create_group('model')
        torchutil.save_model_to_hdf5_group(self._model, h5file['model'])

    def save(self, file_path):
        torch.save({
            'model_state_dict': self._model.state_dict(),
            'encoder_name': self._encoder.name(),
            'encoder_board_width': self._encoder.board_width,
            'encoder_board_height': self._encoder.board_height,
            'temperature': self._temperature,
        }, file_path)

    @classmethod
    def load(cls, file_path):
        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)

        checkpoint = torch.load(file_path)
        model = TheModelClass()  # TheModelClass should be replaced with your model class
        model.load_state_dict(checkpoint['model_state_dict'])
        encoder_name = checkpoint['encoder_name']
        board_width = checkpoint['encoder_board_width']
        board_height = checkpoint['encoder_board_height']
        encoder = encoders.get_encoder_by_name(
            encoder_name,
            (board_width, board_height))
        return cls(model, encoder)