import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

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

class QAgent(Agent):
    def __init__(self, model, encoder, policy='eps-greedy'):
        Agent.__init__(self)
        self.model = model
        self.encoder = encoder
        self.collector = None
        self.temperature = 0.0
        self.policy = policy

        self.last_move_value = 0

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

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

    def set_policy(self, policy):
        if policy not in ('eps-greedy', 'weighted'):
            raise ValueError(policy)
        self.policy = policy

    def select_move(self, game_state):
        board_tensor = self.encoder.encode(game_state)

        # Loop over all legal moves.
        moves = []
        board_tensors = []
        for move in game_state.legal_moves():
            if not move.is_play:
                continue
            moves.append(self.encoder.encode_point(move.point))
            board_tensors.append(board_tensor)
        if not moves:
            return goboard.Move.pass_turn()

        num_moves = len(moves)
        board_tensors = torch.tensor(board_tensors)
        move_vectors = torch.zeros((num_moves, self.encoder.num_points()))
        for i, move in enumerate(moves):
            move_vectors[i][move] = 1

        values = self.model([board_tensors, move_vectors])
        values = values.reshape(len(moves))

        if self.policy == 'eps-greedy':
            ranked_moves = self.rank_moves_eps_greedy(values)
        elif self.policy == 'weighted':
            ranked_moves = self.rank_moves_weighted(values)
        else:
            ranked_moves = None

        for move_idx in ranked_moves:
            point = self.encoder.decode_point_index(moves[move_idx])
            if 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=moves[move_idx],
                    )
                self.last_move_value = float(values[move_idx])
                return goboard.Move.play(point)
        # No legal, non-self-destructive moves less.
        return goboard.Move.pass_turn()

    def rank_moves_eps_greedy(self, values):
        if torch.rand(1).item() < self.temperature:
            values = torch.rand_like(values)
        # This ranks the moves from worst to best.
        ranked_moves = torch.argsort(values)
        # Return them in best-to-worst order.
        return ranked_moves[::-1]

    def rank_moves_weighted(self, values):
        p = values / torch.sum(values)
        p = torch.pow(p, 1.0 / self.temperature)
        p = p / torch.sum(p)
        return Categorical(p).sample().item()

    def train(self, experience, lr=0.1, batch_size=128):
        opt = optim.SGD(self.model.parameters(), lr=lr)
        criterion = nn.MSELoss()

        n = experience.states.shape[0]
        num_moves = self.encoder.num_points()
        y = torch.zeros((n,))
        actions = torch.zeros((n, num_moves))
        for i in range(n):
            action = experience.actions[i]
            reward = experience.rewards[i]
            actions[i][action] = 1
            y[i] = 1 if reward > 0 else 0

        self.model.train()
        opt.zero_grad()
        pred = self.model([experience.states, actions])
        loss = criterion(pred, y)
        loss.backward()
        opt.step()

    def diagnostics(self):
        return {'value': self.last_move_value}

    def save(self, file_path):
        torch.save({
            'model': self.model.state_dict(),
            'encoder': {
                'name': self.encoder.name(),
                'board_width': self.encoder.board_width,
                'board_height': self.encoder.board_height,
            }
        }, file_path)

    @classmethod
    def load(cls, file_path):
        checkpoint = torch.load(file_path)

        # Assuming `model` is a PyTorch model here and `encoder` is a DLGo encoder
        model = Model()  # The Model class should be defined appropriately
        model.load_state_dict(checkpoint['model'])

        encoder_info = checkpoint['encoder']
        encoder = encoders.get_encoder_by_name(
            encoder_info['name'],
            (encoder_info['board_width'], encoder_info['board_height']))

        return cls(model, encoder)


# def load_q_agent(file_path):
#     # 我假设 Model 类是预先定义好的 PyTorch 模型，你需要根据实际情况替换或实现这个模型。
#     checkpoint = torch.load(file_path)
#
#     # Assuming `model` is a PyTorch model here and `encoder` is a DLGo encoder
#     model = Model()  # The Model class should be defined appropriately
#     model.load_state_dict(checkpoint['model'])
#
#     encoder_info = checkpoint['encoder']
#     encoder = encoders.get_encoder_by_name(
#         encoder_info['name'],
#         (encoder_info['board_width'], encoder_info['board_height']))
#
#     return QAgent(model, encoder)