# coding:utf-8
import json
import os
import time
import traceback

import torch
import torch.nn.functional as F
from torch import nn, optim, cuda
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader

from .alpha_zero_mcts import AlphaZeroMCTS
from .chess_board import ChessBoard
from .policy_value_net import PolicyValueNet
from .self_play_dataset import SelfPlayData, SelfPlayDataSet


def exception_handler(train_func):
    """ å¼‚å¸¸å¤„ç�†è£…é¥°å™¨ """
    def wrapper(train_pipe_line, *args, **kwargs):
        try:
            train_func(train_pipe_line)
        except BaseException as e:
            if not isinstance(e, KeyboardInterrupt):
                traceback.print_exc()

            t = time.strftime('%Y-%m-%d_%H-%M-%S',
                              time.localtime(time.time()))
            train_pipe_line.save_model(
                f'last_policy_value_net_{t}.pth', 'train_losses', 'games')

    return wrapper


class PolicyValueLoss(nn.Module):
    """ æ ¹æ�® self-play äº§ç”Ÿçš„ `z` å’Œ `Ï€` è®¡ç®—è¯¯å·® """

    def __init__(self):
        super().__init__()

    def forward(self, p_hat, pi, value, z):
        """ å‰�é¦ˆ

        Parameters
        ----------
        p_hat: Tensor of shape (N, board_len^2)
            å¯¹æ•°åŠ¨ä½œæ¦‚ç�‡å�‘é‡�

        pi: Tensor of shape (N, board_len^2)
            `mcts` äº§ç”Ÿçš„åŠ¨ä½œæ¦‚ç�‡å�‘é‡�

        value: Tensor of shape (N, )
            å¯¹æ¯�ä¸ªå±€é�¢çš„ä¼°å€¼

        z: Tensor of shape (N, )
            æœ€ç»ˆçš„æ¸¸æˆ�ç»“æ�œç›¸å¯¹æ¯�ä¸€ä¸ªç�©å®¶çš„å¥–èµ�
        """
        value_loss = F.mse_loss(value, z)
        policy_loss = -torch.sum(pi*p_hat, dim=1).mean()
        loss = value_loss + policy_loss
        return loss


class TrainModel:
    """ è®­ç»ƒæ¨¡å�‹ """

    def __init__(self, board_len=9, lr=0.01, n_self_plays=1500, n_mcts_iters=500,
                 n_feature_planes=4, batch_size=500, start_train_size=500, check_frequency=100,
                 n_test_games=10, c_puct=4, is_use_gpu=True, is_save_game=False, **kwargs):
        """
        Parameters
        ----------
        board_len: int
            æ£‹ç›˜å¤§å°�

        lr: float
            å­¦ä¹ ç�‡

        n_self_plays: int
            è‡ªæˆ‘å�šå¼ˆæ¸¸æˆ�å±€æ•°

        n_mcts_iters: int
            è’™ç‰¹å�¡æ´›æ ‘æ�œç´¢æ¬¡æ•°

        n_feature_planes: int
            ç‰¹å¾�å¹³é�¢ä¸ªæ•°

        batch_size: int
            mini-batch çš„å¤§å°�

        start_train_size: int
            å¼€å§‹è®­ç»ƒæ¨¡å�‹æ—¶çš„æœ€å°�æ•°æ�®é›†å°ºå¯¸

        check_frequency: int
            æµ‹è¯•æ¨¡å�‹çš„é¢‘ç�‡

        n_test_games: int
            æµ‹è¯•æ¨¡å�‹æ—¶ä¸�å�†å�²æœ€ä¼˜æ¨¡å�‹çš„æ¯”èµ›å±€æ•°

        c_puct: float
            æ�¢ç´¢å¸¸æ•°

        is_use_gpu: bool
            æ˜¯å�¦ä½¿ç”¨ GPU

        is_save_game: bool
            æ˜¯å�¦ä¿�å­˜è‡ªå¯¹å¼ˆçš„æ£‹è°±
        """
        self.c_puct = c_puct
        self.is_use_gpu = is_use_gpu
        self.batch_size = batch_size
        self.n_self_plays = n_self_plays
        self.n_test_games = n_test_games
        self.n_mcts_iters = n_mcts_iters
        self.is_save_game = is_save_game
        self.check_frequency = check_frequency
        self.start_train_size = start_train_size
        self.device = torch.device(
            'cuda:0' if is_use_gpu and cuda.is_available() else 'cpu')
        self.chess_board = ChessBoard(board_len, n_feature_planes)

        # åˆ›å»ºç­–ç•¥-ä»·å€¼ç½‘ç»œå’Œè’™ç‰¹å�¡æ´›æ�œç´¢æ ‘
        self.policy_value_net = self.__get_policy_value_net(board_len)
        self.mcts = AlphaZeroMCTS(
            self.policy_value_net, c_puct=c_puct, n_iters=n_mcts_iters, is_self_play=True)

        # åˆ›å»ºä¼˜åŒ–å™¨å’Œæ�Ÿå¤±å‡½æ•°
        self.optimizer = optim.Adam(
            self.policy_value_net.parameters(), lr=lr, weight_decay=1e-4)
        self.criterion = PolicyValueLoss()
        self.lr_scheduler = MultiStepLR(
            self.optimizer, [1500, 2500], gamma=0.1)

        # åˆ›å»ºæ•°æ�®é›†
        self.dataset = SelfPlayDataSet(board_len)

        # è®°å½•æ•°æ�®
        self.train_losses = self.__load_data('log/train_losses.json')
        self.games = self.__load_data('log/games.json')

    def __self_play(self):
        """ è‡ªæˆ‘å�šå¼ˆä¸€å±€

        Returns
        -------
        self_play_data: namedtuple
            è‡ªæˆ‘å�šå¼ˆæ•°æ�®ï¼Œæœ‰ä»¥ä¸‹ä¸‰ä¸ªæˆ�å‘˜:
            * `pi_list`: è’™ç‰¹å�¡æ´›æ ‘æ�œç´¢äº§ç”Ÿçš„åŠ¨ä½œæ¦‚ç�‡å�‘é‡� Ï€ ç»„æˆ�çš„åˆ—è¡¨
            * `z_list`: ä¸€å±€ä¹‹ä¸­æ¯�ä¸ªåŠ¨ä½œçš„ç�©å®¶ç›¸å¯¹æœ€å��çš„æ¸¸æˆ�ç»“æ�œçš„å¥–èµ�åˆ—è¡¨
            * `feature_planes_list`: ä¸€å±€ä¹‹ä¸­æ¯�ä¸ªåŠ¨ä½œå¯¹åº”çš„ç‰¹å¾�å¹³é�¢ç»„æˆ�çš„åˆ—è¡¨
        """
        # åˆ�å§‹åŒ–æ£‹ç›˜å’Œæ•°æ�®å®¹å™¨
        self.policy_value_net.eval()
        self.chess_board.clear_board()
        pi_list, feature_planes_list, players = [], [], []
        action_list = []

        # å¼€å§‹ä¸€å±€æ¸¸æˆ�
        while True:
            action, pi = self.mcts.get_action(self.chess_board)

            # ä¿�å­˜æ¯�ä¸€æ­¥çš„æ•°æ�®
            feature_planes_list.append(self.chess_board.get_feature_planes())
            players.append(self.chess_board.current_player)
            action_list.append(action)
            pi_list.append(pi)
            self.chess_board.do_action(action)

            # åˆ¤æ–­æ¸¸æˆ�æ˜¯å�¦ç»“æ�Ÿ
            is_over, winner = self.chess_board.is_game_over()
            if is_over:
                if winner is not None:
                    z_list = [1 if i == winner else -1 for i in players]
                else:
                    z_list = [0]*len(players)
                break

        # é‡�ç½®æ ¹èŠ‚ç‚¹
        self.mcts.reset_root()

        # è¿”å›�æ•°æ�®
        if self.is_save_game:
            self.games.append(action_list)

        self_play_data = SelfPlayData(
            pi_list=pi_list, z_list=z_list, feature_planes_list=feature_planes_list)
        return self_play_data

    @exception_handler
    def train(self):
        """ è®­ç»ƒæ¨¡å�‹ """
        for i in range(self.n_self_plays):
            print(f'ğŸ�¹ æ­£åœ¨è¿›è¡Œç¬¬ {i+1} å±€è‡ªæˆ‘å�šå¼ˆæ¸¸æˆ�...')
            self.dataset.append(self.__self_play())

            # å¦‚æ�œæ•°æ�®é›†ä¸­çš„æ•°æ�®é‡�å¤§äº� start_train_size å°±è¿›è¡Œä¸€æ¬¡è®­ç»ƒ
            if len(self.dataset) >= self.start_train_size:
                data_loader = iter(DataLoader(
                    self.dataset, self.batch_size, shuffle=True, drop_last=False))
                print('ğŸ’Š å¼€å§‹è®­ç»ƒ...')

                self.policy_value_net.train()
                # éš�æœºé€‰å‡ºä¸€æ‰¹æ•°æ�®æ�¥è®­ç»ƒï¼Œé˜²æ­¢è¿‡æ‹Ÿå�ˆ
                feature_planes, pi, z = next(data_loader)
                feature_planes = feature_planes.to(self.device)
                pi, z = pi.to(self.device), z.to(self.device)
                for _ in range(5):
                    # å‰�é¦ˆ
                    p_hat, value = self.policy_value_net(feature_planes)
                    # æ¢¯åº¦æ¸…é›¶
                    self.optimizer.zero_grad()
                    # è®¡ç®—æ�Ÿå¤±
                    loss = self.criterion(p_hat, pi, value.flatten(), z)
                    # è¯¯å·®å��å�‘ä¼ æ’­
                    loss.backward()
                    # æ›´æ–°å�‚æ•°
                    self.optimizer.step()
                    # å­¦ä¹ ç�‡é€€ç�«
                    self.lr_scheduler.step()

                # è®°å½•è¯¯å·®
                self.train_losses.append([i, loss.item()])
                print(f"ğŸš© train_loss = {loss.item():<10.5f}\n")

            # æµ‹è¯•æ¨¡å�‹
            if (i+1) % self.check_frequency == 0:
                self.__test_model()

    def __test_model(self):
        """ æµ‹è¯•æ¨¡å�‹ """
        os.makedirs('model', exist_ok=True)

        model_path = 'model/best_policy_value_net.pth'

        # å¦‚æ�œæœ€ä½³æ¨¡å�‹ä¸�å­˜åœ¨ä¿�å­˜å½“å‰�æ¨¡å�‹ä¸ºæœ€ä½³æ¨¡å�‹
        if not os.path.exists(model_path):
            torch.save(self.policy_value_net, model_path)
            return

        # è½½å…¥å�†å�²æœ€ä¼˜æ¨¡å�‹
        best_model = torch.load(model_path)  # type:PolicyValueNet
        best_model.eval()
        best_model.set_device(self.is_use_gpu)
        mcts = AlphaZeroMCTS(best_model, self.c_puct, self.n_mcts_iters)
        self.mcts.set_self_play(False)
        self.policy_value_net.eval()

        # å¼€å§‹æ¯”èµ›
        print('ğŸ©º æ­£åœ¨æµ‹è¯•å½“å‰�æ¨¡å�‹...')
        n_wins = 0
        for i in range(self.n_test_games):
            self.chess_board.clear_board()
            self.mcts.reset_root()
            mcts.reset_root()
            while True:
                # å½“å‰�æ¨¡å�‹èµ°ä¸€æ­¥
                is_over, winner = self.__do_mcts_action(self.mcts)
                if is_over:
                    n_wins += int(winner == ChessBoard.BLACK)
                    break
                # å�†å�²æœ€ä¼˜æ¨¡å�‹èµ°ä¸€æ­¥
                is_over, winner = self.__do_mcts_action(mcts)
                if is_over:
                    break

        # å¦‚æ�œèƒœç�‡å¤§äº� 55%ï¼Œå°±ä¿�å­˜å½“å‰�æ¨¡å�‹ä¸ºæœ€ä¼˜æ¨¡å�‹
        win_prob = n_wins/self.n_test_games
        if win_prob > 0.55:
            torch.save(self.mcts.policy_value_net, model_path)
            print(f'ğŸ¥‡ ä¿�å­˜å½“å‰�æ¨¡å�‹ä¸ºæœ€ä¼˜æ¨¡å�‹ï¼Œå½“å‰�æ¨¡å�‹èƒœç�‡ä¸ºï¼š{win_prob:.1%}\n')
        else:
            print(f'ğŸ�ƒ ä¿�æŒ�å�†å�²æœ€ä¼˜æ¨¡å�‹ä¸�å�˜ï¼Œå½“å‰�æ¨¡å�‹èƒœç�‡ä¸ºï¼š{win_prob:.1%}\n')

        self.mcts.set_self_play(True)

    def save_model(self, model_name: str, loss_name: str, game_name: str):
        """ ä¿�å­˜æ¨¡å�‹

        Parameters
        ----------
        model_name: str
            æ¨¡å�‹æ–‡ä»¶å��ç§°ï¼Œä¸�åŒ…å�«å��ç¼€

        loss_name: str
            æ�Ÿå¤±æ–‡ä»¶å��ç§°ï¼Œä¸�åŒ…å�«å��ç¼€

        game_name: str
            è‡ªå¯¹å¼ˆæ£‹è°±å��ç§°ï¼Œä¸�åŒ…å�«å��ç¼€
        """
        os.makedirs('model', exist_ok=True)

        path = f'model/{model_name}.pth'
        self.policy_value_net.eval()
        torch.save(self.policy_value_net, path)
        print(f'ğŸ�‰ å·²å°†å½“å‰�æ¨¡å�‹ä¿�å­˜åˆ° {os.path.join(os.getcwd(), path)}')

        # ä¿�å­˜æ•°æ�®
        with open(f'log/{loss_name}.json', 'w', encoding='utf-8') as f:
            json.dump(self.train_losses, f)

        if self.is_save_game:
            with open(f'log/{game_name}.json', 'w', encoding='utf-8') as f:
                json.dump(self.games, f)


    def __do_mcts_action(self, mcts):
        """ è�·å�–åŠ¨ä½œ """
        action = mcts.get_action(self.chess_board)
        self.chess_board.do_action(action)
        is_over, winner = self.chess_board.is_game_over()
        return is_over, winner

    def __get_policy_value_net(self, board_len=9):
        """ åˆ›å»ºç­–ç•¥-ä»·å€¼ç½‘ç»œï¼Œå¦‚æ�œå­˜åœ¨å�†å�²æœ€ä¼˜æ¨¡å�‹åˆ™ç›´æ�¥è½½å…¥æœ€ä¼˜æ¨¡å�‹ """
        os.makedirs('model', exist_ok=True)

        best_model = 'best_policy_value_net.pth'
        history_models = sorted(
            [i for i in os.listdir('model') if i.startswith('last')])

        # ä»�å�†å�²æ¨¡å�‹ä¸­é€‰å�–æœ€æ–°æ¨¡å�‹
        model = history_models[-1] if history_models else best_model
        model = f'model/{model}'
        if os.path.exists(model):
            print(f'ğŸ’� è½½å…¥æ¨¡å�‹ {model} ...\n')
            net = torch.load(model).to(self.device)  # type:PolicyValueNet
            net.set_device(self.is_use_gpu)
        else:
            net = PolicyValueNet(n_feature_planes=self.chess_board.n_feature_planes,
                                 is_use_gpu=self.is_use_gpu, board_len=board_len).to(self.device)

        return net

    def __load_data(self, path: str):
        """ è½½å…¥å�†å�²æ�Ÿå¤±æ•°æ�® """
        data = []
        try:
            with open(path, encoding='utf-8') as f:
                data = json.load(f)
        except:
            os.makedirs('log', exist_ok=True)

        return data
