# -*- coding: utf-8 -*-

import numpy as np
import pygame


class Board:
    """
    board: 0: available, 1: player1, 2: player2
    """
    board: np.ndarray

    current_player = 1
    last_move = -1
    move_count = 0

    def __init__(self, width=8, height=8, n_in_rows=4):
        self.width = width
        self.height = height
        self.n_in_rows = n_in_rows
        if width < self.n_in_rows or height < self.n_in_rows:
            raise ValueError("Board size must be greater than or equal to n_in_rows")

    def init_board(self, start_player=1):
        self.board = np.zeros((self.height, self.width))
        self.current_player = start_player
        self.move_count = 0

    def position2move(self, pos):
        """

        :param pos: (y, x)
        :return: move x + w * y
        """
        assert 0 <= pos[1] < self.width and 0 <= pos[0] < self.height, f'Position {pos} out of bounds'
        return pos[1] + self.width * pos[0]

    def move2position(self, move: int):
        """
        :param move:
        :return: pos: (y, x)
        """
        assert 0 <= move < self.height * self.width, f'Move {move} out of bounds'
        return move // self.width, move % self.width

    def game_over(self):
        if self.last_move == -1:
            return False, None
        if self.move_count < self.n_in_rows * 2 - 1:
            return False, None
        if self.move_count == self.width * self.height:
            return True, None

        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        last_pos = self.move2position(self.last_move)

        for d in directions:
            count = 1  # 当前刚落子的棋子
            # 正向检查
            y, x = last_pos[0] + d[0], last_pos[1] + d[1]
            while 0 <= x < self.width and 0 <= y < self.height and self.board[y, x] == self.opponent_player:
                count += 1
                x += d[1]
                y += d[0]

            # 反向检查
            y, x = last_pos[0] - d[0], last_pos[1] - d[1]
            while 0 <= x < self.width and 0 <= y < self.height and self.board[y, x] == self.opponent_player:
                count += 1
                x -= d[1]
                y -= d[0]

            if count >= self.n_in_rows:
                return True, self.opponent_player
        return False, None

    @property
    def state(self):
        state = np.zeros((4, self.height, self.width))
        state[0] = self.board == self.current_player
        state[1] = self.board == self.opponent_player
        if self.last_move != -1:
            last_pos = self.move2position(self.last_move)
            state[2, last_pos[0], last_pos[1]] = 1.0
        if self.move_count % 2 == 0:
            state[3][:, :] = 1.0
        # https://github.com/junxiaosong/AlphaZero_Gomoku/issues/33
        # original state[:, ::-1, :]
        return state

    @property
    def opponent_player(self):
        return 1 if (self.current_player == 2) else 2

    def get_available_moves(self):
        rows, columns = np.where(self.board == 0)
        return rows * self.width + columns

class TreeNode(object):

    def __init__(self, parent, prior_p):
        self.parent = parent
        self.children = {}

        self.n_visits = 0
        self.Q = 0
        self.u = 0

        self.P = prior_p

    def expand(self, action_priors):
        """Expand tree by creating new children.
        :param action_priors: a list of tuples of actions and their prior probability
            according to the policy function.
        :return:
        """

        for action, prior_prob in action_priors:
            if action not in self.children:
                self.children[action] = TreeNode(self, prior_prob)

    def select(self, c_puct):
        """Select action among children that gives maximum action value Q plus bonus u(P).
        :return: A tuple of (action, next_node)
        """

        return max(self.children.items(), key=lambda act_node: act_node[1].get_value(c_puct))

    def update(self, leaf_value):
        """Update node values from leaf evaluation.
        :param leaf_value: the value of subtree evaluation from the current player's
            perspective.
        """

        # Count visit.
        self.n_visits += 1
        # Update Q, a running average of values for all visits.
        # (n * Q + v) / (n + 1) - Q = (v - Q) / (n + 1)
        self.Q += 1.0 * (leaf_value - self.Q) / self.n_visits

    def update_recursive(self, leaf_value):
        """Like a call to update(), but applied recursively for all ancestors."""

        # If it is not root, this node's parent should be updated first.
        if self.parent:
            self.parent.update_recursive(-leaf_value)
        self.update(leaf_value)

    def get_value(self, c_puct):
        """Calculate and return the value for this node.
        It is a combination of leaf evaluations Q, and this node's prior
        adjusted for its visit count, u.
        :param c_puct: a number in (0, inf) controlling the relative impact of
            value Q, and prior probability P, on this node's score.
        """

        self.u = (c_puct * self.P * np.sqrt(self.parent.n_visits) / (1 + self.n_visits))
        return self.Q + self.u

    def is_leaf(self):
        """Check if leaf node (i.e. no nodes below this have been expanded)."""

        return self.children == {}

    def is_root(self):
        return self.parent is None


class MCTSPure(object):
    def __init__(self, policy_value_fn, c_puct=5, n_playout=10000):
        """
        :param policy_value_fn: a function that takes in a board state and outputs
            a list of (action, probability) tuples and also a score in [-1, 1]
            (i.e. the expected value of the end game score from the current
            player's perspective) for the current player.
        :param c_puct: a number in (0, inf) that controls how quickly exploration
            converges to the maximum-value policy. A higher value means
            relying on the prior more.
        :param n_playout:
        """
        self.root = TreeNode(None, 1.0)
        self.policy_value_fn = policy_value_fn
        self.c_puct = c_puct
        self.n_playout = n_playout

    def playout(self, state: Board):
        """Run a single playout from the root to the leaf, getting a value at
        the leaf and propagating it back through its parents.
        State is modified in-place, so a copy must be provided.
        """
        node = self.root
        while not node.is_leaf():
            action, node = node.select(self.c_puct)
            state.perform_move(action)
        action_probs, _ = self.policy_value_fn(state)

        end, winner = state.game_over()
        if not end:
            node.expand(action_probs)

        # Evaluate the leaf node by random rollout
        leaf_value = self.evaluate_rollout(state)
        # Update value and visit count of nodes in this traversal.
        """
        Q:  Why is the negative leaf_value for update_recursive function?
        A:  We use the negative value of the state, this is because alternate levels 
            in the search tree are from the perspective of different players and 
            the Q-values are in fact used by the parent node in select stage.
        Reference https://github.com/junxiaosong/AlphaZero_Gomoku/issues/25
        """
        node.update_recursive(-leaf_value)

    def evaluate_rollout(self, state: Board, limit=1000):
        """Use the rollout policy to play until the end of the game,
        returning +1 if the current player wins, -1 if the opponent wins,
        and 0 if it is a tie."""
        assert limit >= 1, 'Limit must be > 0'
        player = state.current_player
        winner = None
        for i in range(limit):
            end, winner = state.game_over()
            if end:
                break
            action_probs = random_rollout_policy_fn(state)
            max_action = max(action_probs, key=itemgetter(1))[0]
            state.perform_move(max_action)
        else:
            # If no break from the loop, issue a warning.
            print("WARNING: rollout reached move limit")
        if winner is None:  # tie
            return 0
        else:
            return 1 if winner == player else -1

    def get_move(self, state):
        """Runs all playouts sequentially and returns the most visited action.
        state: the current game state
        :return: the selected action
        """
        for n in range(self.n_playout):
            state_copy = deepcopy(state)
            self.playout(state_copy)
        return max(self.root.children.items(), key=lambda act_node: act_node[1].n_visits)[0]

    def update_with_move(self, last_move):
        """Step forward in the tree, keeping everything we already know
        about the subtree.
        """
        if last_move in self.root.children:
            self.root = self.root.children[last_move]
            self.root.parent = None
        else:
            self.root = TreeNode(None, 1.0)

def softmax(x):
    probs = np.exp(x - np.max(x))
    probs /= np.sum(probs)
    return probs

class MCTSAlphaZero(MCTSPure):
    def playout(self, state: Board):
        """Run a single playout from the root to the leaf, getting a value at
        the leaf and propagating it back through its parents.
        State is modified in-place, so a copy must be provided.
        """
        node = self.root
        while not node.is_leaf():
            action, node = node.select(self.c_puct)
            state.perform_move(action)
        # Evaluate the leaf using a network which outputs a list of
        # (action, probability) tuples player and also a score v in [-1, 1]
        # for the current player.
        action_probs, leaf_value = self.policy_value_fn(state)
        end, winner = state.game_over()
        if not end:
            node.expand(action_probs)
        else:
            # for end state，return the "true" leaf_value
            if winner is None:
                leaf_value = 0
            elif winner == state.current_player:
                leaf_value = 1
            else:
                leaf_value = -1
        # Update value and visit count of nodes in this traversal.
        node.update_recursive(-leaf_value)

    def get_move_probs(self, state, temperature=1e-3):
        """Run all playouts sequentially and return the available actions and
        their corresponding probabilities.
        state: the current game state
        temperature: temperature parameter in (0, 1] controls the level of exploration
        """
        for n in range(self.n_playout):
            state_copy = deepcopy(state)
            self.playout(state_copy)

        # calc the move probabilities based on visit counts at the root node
        act_visits = [(act, node.n_visits) for act, node in self.root.children.items()]
        acts, visits = zip(*act_visits)
        act_probs = softmax(1.0 / temperature * np.log(np.array(visits) + 1e-10))
        return acts, act_probs

def random_rollout_policy_fn(board: Board):
    """a coarse, fast version of policy_fn used in the rollout phase."""
    # rollout randomly
    available_moves = board.get_available_moves()
    action_probs = np.random.rand(len(available_moves))
    return zip(available_moves, action_probs)


def random_policy_value_fn(board: Board):
    """a function that takes in a state and outputs a list of (action, probability)
    tuples and a score for the state"""
    # return uniform probabilities and 0 score for pure MCTSPure
    available_moves = board.get_available_moves()
    action_probs = np.ones(len(available_moves)) / len(available_moves)
    return zip(available_moves, action_probs), 0

class BaceMCTSPlayer(object):
    """AI player based on MCTS"""

    def __init__(self, mcts):
        self.player = None
        self.mcts = mcts

    def set_player(self, player: int):
        self.player = player

    def reset_player(self):
        self.mcts.update_with_move(-1)

    def act(self, board: Board):
        raise NotImplementedError()


class MCTSPlayer(BaceMCTSPlayer):
    """AI player based on MCTSPure"""

    def __init__(self, c_puct=5, n_playout=2000):
        super().__init__(MCTSPure(random_policy_value_fn, c_puct, n_playout))

    def act(self, board: Board):
        sensible_moves = board.get_available_moves()
        if len(sensible_moves) > 0:
            move = self.mcts.get_move(board)
            # https://github.com/junxiaosong/AlphaZero_Gomoku/issues/108
            self.mcts.update_with_move(-1)
            # self.mcts.update_with_move(move)
            return move
        else:
            print("WARNING: the board is full")


class MCTSAlphaZeroPlayer(BaceMCTSPlayer):
    """AI player based on MCTSAlphaZero"""

    def __init__(self, policy_value_function, c_puct=5, n_playout=2000, dirichlet=0.3, selfplay=False):
        super().__init__(MCTSAlphaZero(policy_value_function, c_puct, n_playout))
        self.dirichlet = dirichlet
        self.selfplay = selfplay

    def act(self, board, temperature=1e-3, return_prob=False):
        sensible_moves = board.get_available_moves()
        # the pi vector returned by MCTS as in the alphaGo Zero paper
        move_probs = np.zeros(board.width * board.height)
        if len(sensible_moves) > 0:
            acts, probs = self.mcts.get_move_probs(board, temperature)
            move_probs[list(acts)] = probs
            if self.selfplay:
                # add Dirichlet Noise for exploration (needed for
                # self-play training)
                move = np.random.choice(
                    acts,
                    p=0.75 * probs + 0.25 * np.random.dirichlet(self.dirichlet * np.ones(len(probs)))
                )
                # update the root node and reuse the search tree
                self.mcts.update_with_move(move)
            else:
                # with the default temperature=1e-3, it is almost equivalent
                # to choosing the move with the highest prob
                move = np.random.choice(acts, p=probs)
                # reset the root node
                self.mcts.update_with_move(-1)
                # location = board.move_to_location(move)
                # print("AI move: %d,%d\n" % (location[0], location[1]))

            if return_prob:
                return move, move_probs
            else:
                return move
        else:
            print("WARNING: the board is full")

# -*- coding: utf-8 -*-
import os
import warnings
from pathlib import Path

import numpy as np
import paddle
from paddle import nn
from paddle.static import InputSpec

from game import Board


class ResidualBlock(nn.Layer):
    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.Conv2D(channels, channels, kernel_size=3, padding=1, bias_attr=False),
            nn.BatchNorm2D(channels),
            nn.ReLU(),
            nn.Conv2D(channels, channels, kernel_size=3, padding=1, bias_attr=False),
            nn.BatchNorm2D(channels)
        )

    def forward(self, x):
        residual = x  # 旁路连接
        out = self.conv_block(x)
        out += residual  # 跳跃连接
        return nn.functional.relu(out)  # 经过 ReLU 激活


class Net(nn.Layer):
    """
    https://github.com/junxiaosong/AlphaZero_Gomoku
    """

    def __init__(self, board_width, board_height, num_blocks=5):
        super().__init__()

        self.board_width = board_width
        self.board_height = board_height

        # common layers
        self.common_layers = nn.Sequential(
            nn.Conv2D(4, 128, kernel_size=3, stride=1, padding=2),
            nn.BatchNorm2D(128),
            *[ResidualBlock(128) for _ in range(num_blocks)]
        )

        # action policy layers
        self.act_layers = nn.Sequential(
            nn.Conv2D(128, 4, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(4 * (board_height + 2) * (board_width + 2), board_width * board_height),
            nn.LogSoftmax()
        )

        # action value layers
        self.val_layers = nn.Sequential(
            nn.Conv2D(128, 2, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(2 * (board_height + 2) * (board_width + 2), 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.common_layers(x)
        x_act = self.act_layers(x)
        x_val = self.val_layers(x)
        return x_act, x_val


class PolicyValueNet(object):
    def __init__(self, board_width, board_height, model_file=None):
        self.board_width = board_width
        self.board_height = board_height
        self.model_file = model_file

        self.l2_const = 1e-4  # coef of l2 penalty

        self.net = Net(self.board_width, self.board_height)

        self.optimizer = paddle.optimizer.Adam(parameters=self.net.parameters(), weight_decay=self.l2_const,
                                               learning_rate=2e-3)

        if model_file is not None:
            # if not os.path.isfile(model_file):
            #     raise FileNotFoundError(f'{model_file} is not exist')
            net_params = paddle.load(model_file + '.pdparams')
            self.net.set_state_dict(net_params)
            opt_params = paddle.load(model_file + '.pdopt')
            self.optimizer.set_state_dict(opt_params)

    def policy_value(self, state: np.ndarray | list):
        """
        input: a batch of states
        output: a batch of action probabilities and state values
        """
        log_act_probs, value = self.net(paddle.to_tensor(state, 'float32'))
        act_probs = log_act_probs.exp()
        return act_probs.numpy(), value.numpy()

    def policy_value_fn(self, board: Board):
        """
        input: board
        output: a list of (action, probability) tuples for each available
        action and the score of the board state
        """
        legal_positions = board.get_available_moves()
        current_state = np.ascontiguousarray(board.state.reshape(-1, 4, self.board_height, self.board_width))
        act_probs, value = self.policy_value(current_state)
        value = value[0][0]
        act_probs = zip(legal_positions, act_probs[0][legal_positions])
        return act_probs, value

    def train_step(self, state_batch, mcts_probs, winner_batch, lr):
        """perform a training step"""
        state_batch = paddle.to_tensor(state_batch, dtype='float32')
        mcts_probs = paddle.to_tensor(mcts_probs, dtype='float32')
        winner_batch = paddle.to_tensor(winner_batch, dtype='float32')

        # zero the parameter gradients
        self.optimizer.clear_grad()
        # set learning rate
        self.optimizer.set_lr(lr)

        # forward
        log_act_probs, value = self.net(state_batch)
        # define the loss = (z - v)^2 - pi^T * log(p) + c||theta||^2
        # Note: the L2 penalty is incorporated in optimizer
        value = paddle.reshape(x=value, shape=[-1])
        value_loss = nn.functional.mse_loss(value, winner_batch)
        policy_loss = -paddle.mean(paddle.sum(mcts_probs * log_act_probs, 1))
        loss = value_loss + policy_loss
        # backward and optimize
        loss.backward()
        self.optimizer.step()
        # calc policy entropy, for monitoring only
        entropy = -paddle.mean(
            paddle.sum(paddle.exp(log_act_probs) * log_act_probs, 1)
        )
        return loss.item(), entropy.item()

    def get_policy_param(self):
        return self.net.state_dict()

    def set_policy_param(self, param):
        self.net.set_state_dict(param)

    def save_model(self, path: Path | str):
        """ save model params to file """
        if isinstance(path, Path):
            path = str(path)
        net_params = self.get_policy_param()  # get model params
        paddle.save(net_params, path + '.pdparams')
        opt_params = self.optimizer.state_dict()
        paddle.save(opt_params, path + '.pdopt')

    def save_for_predict(self, path: Path | str):
        if isinstance(path, Path):
            path = str(path)
        # save inferencing format model
        paddle.jit.save(
            self.net, path, input_spec=[
                InputSpec(shape=[None, 4, self.board_height, self.board_width], dtype='float32')
            ]
        )

# -*- coding: utf-8 -*-
import os
import random
from collections import deque
from pathlib import Path

import numpy as np

from game import Board, GameGomuku
from mcts import MCTSPlayer, MCTSAlphaZeroPlayer
from policy_value_net import PolicyValueNet

MODEL_DIR = Path('./models')


class TrainPipeline:
    def __init__(self, name, init_model=None, **kwargs):
        self.name = name

        # params of the board and the game
        self.board_width = kwargs.get('board_width', 8)
        self.board_height = kwargs.get('board_height', 8)
        self.n_in_rows = kwargs.get('n_in_rows', 4)
        self.board = Board(width=self.board_width,
                           height=self.board_height,
                           n_in_rows=self.n_in_rows)
        self.game = GameGomuku(self.board)

        # training params
        self.learning_rate = kwargs.get('learning_rate', 2e-3)
        self.lr_multiplier = kwargs.get('lr_multiplier', 1.0)  # adaptively adjust the learning rate based on KL
        self.temperature = kwargs.get('temperature', 1.0)
        self.n_playout = kwargs.get('n_playout', 400)
        self.c_puct = kwargs.get('c_puct', 5)
        self.buffer_size = kwargs.get('buffer_size', 10000)
        self.batch_size = kwargs.get('batch_size', 512)
        self.data_buffer = deque(maxlen=self.buffer_size)

        self.play_batch_size = 1
        self.epochs = kwargs.get('epochs', 5)  # num of train_steps for each update
        self.kl_targ = kwargs.get('kl_targ', 0.02)
        self.check_freq = kwargs.get('check_freq', 50)
        self.game_batch_num = kwargs.get('game_batch_num', 1500)
        self.dirichlet = kwargs.get('dirichlet', 0.3)

        self.best_win_ratio = 0.0
        # num of simulations used for the pure mcts, which is used as
        # the opponent to evaluate the trained policy
        self.pure_mcts_playout_num = kwargs.get('pure_mcts_playout_num', 1000)
        if init_model is not None:
            logger.info(f'Loading model from [underline]{init_model}[/]')
            # start training from an initial policy-value net
            self.policy_value_net = PolicyValueNet(self.board_width,
                                                   self.board_height,
                                                   model_file=init_model)
        else:
            # start training from a new policy-value net
            self.policy_value_net = PolicyValueNet(self.board_width,
                                                   self.board_height)
        self.mcts_player = MCTSAlphaZeroPlayer(self.policy_value_net.policy_value_fn,
                                               c_puct=self.c_puct,
                                               n_playout=self.n_playout,
                                               dirichlet=self.dirichlet,
                                               selfplay=True)

        self.episode_len = 0

    def get_equi_data(self, play_data):
        """augment the data set by rotation and flipping
        :param play_data: [(state, mcts_prob, winner_z), ..., ...]
        """
        extend_data = []
        for state, mcts_prob, winner in play_data:
            for i in ([0, 1, 2, 3] if (self.board_width == self.board_height) else [0, 2]):
                # rotate counterclockwise
                equi_state = np.rot90(state, i, axes=(1, 2))
                equi_mcts_prob = np.rot90(mcts_prob.reshape(self.board_height, self.board_width), i)
                extend_data.append((equi_state, equi_mcts_prob.flatten(), winner))
                # flip horizontally
                equi_state = equi_state[:, :, ::-1]
                equi_mcts_prob = np.fliplr(equi_mcts_prob)
                extend_data.append((equi_state, equi_mcts_prob.flatten(), winner))
        return extend_data

    def collect_selfplay_data(self, n_games=1):
        """collect self-play data for training"""
        for i in range(n_games):
            winner, play_data = self.game.start_self_play(self.mcts_player, temperature=self.temperature)
            play_data = list(play_data)[:]
            self.episode_len = len(play_data)
            # augment the data
            play_data = self.get_equi_data(play_data)
            self.data_buffer.extend(play_data)

    def policy_update(self):
        """update the policy-value net"""
        mini_batch = random.sample(self.data_buffer, self.batch_size)
        state_batch, mcts_probs_batch, winner_batch = zip(*mini_batch)
        old_probs, old_v = self.policy_value_net.policy_value(state_batch)
        for i in range(self.epochs):
            loss, entropy = self.policy_value_net.train_step(
                state_batch,
                mcts_probs_batch,
                winner_batch,
                self.learning_rate * self.lr_multiplier)
            new_probs, new_v = self.policy_value_net.policy_value(state_batch)
            kl = np.mean(np.sum(old_probs * (np.log(old_probs + 1e-10) - np.log(new_probs + 1e-10)), axis=1))
            if kl > self.kl_targ * 4:  # early stopping if D_KL diverges badly
                break
        # adaptively adjust the learning rate
        if kl > self.kl_targ * 2 and self.lr_multiplier > 0.1:
            self.lr_multiplier /= 1.5
        elif kl < self.kl_targ / 2 and self.lr_multiplier < 20:
            self.lr_multiplier *= 1.5

        explained_var_old = (1 -
                             np.var(np.array(winner_batch) - old_v.flatten()) /
                             np.var(np.array(winner_batch)))
        explained_var_new = (1 -
                             np.var(np.array(winner_batch) - new_v.flatten()) /
                             np.var(np.array(winner_batch)))
        logger.info(f'[bold]Train[/] Step '
                    f'kl: {kl:.5f}, '
                    f'lr_multiplier: {self.lr_multiplier:.3f}, '
                    f'loss: {loss}, '
                    f'entropy: {entropy}, '
                    f'explained_var_old: {explained_var_old:.3f}, '
                    f'explained_var_new: {explained_var_new:.3f}')

    def policy_evaluate(self, n_games=10):
        """
        Evaluate the trained policy by playing against the pure MCTS player
        Note: this is only for monitoring the progress of training
        """
        logger.info(f'[bold]Evaluating[/]')
        current_mcts_player = MCTSAlphaZeroPlayer(self.policy_value_net.policy_value_fn,
                                                  c_puct=self.c_puct,
                                                  n_playout=self.n_playout)
        pure_mcts_player = MCTSPlayer(c_puct=5,
                                      n_playout=self.pure_mcts_playout_num)
        win, lose, tie = 0, 0, 0
        for i in range(n_games):
            winner = self.game.start_play(current_mcts_player,
                                          pure_mcts_player,
                                          start_player=i % 2 + 1, )
            if winner == 1:
                logger.info(f'    {i + 1}/{n_games} [bold]Win[/]')
                win += 1
            elif winner == 2:
                logger.info(f'    {i + 1}/{n_games} [bold]Lose[/]')
                lose += 1
            else:
                logger.info(f'    {i + 1}/{n_games} [bold]Tie[/]')
                tie += 1
        win_ratio = 1.0 * (win + 0.5 * tie) / n_games
        logger.info(f"num_playouts:{self.pure_mcts_playout_num}, "
                    f"win: {win}, lose: {lose}, tie: {tie}")
        return win_ratio

    def run(self):
        """run the training pipeline"""
        try:
            logger.info('Start training.')
            for i in range(self.game_batch_num):
                logger.info(f'[bold]Selfplay[/] batch {i + 1}')
                self.collect_selfplay_data(self.play_batch_size)
                if len(self.data_buffer) > self.batch_size:
                    self.policy_update()
                # check the performance of the current model,
                # and save the model params
                if (i + 1) % self.check_freq == 0:
                    self.policy_value_net.save_model(
                        MODEL_DIR / self.name / 'train' / f'current_step-{i + 1}'
                    )
                    win_ratio = self.policy_evaluate()
                    if win_ratio > self.best_win_ratio:
                        logger.info("[bold]New best policy!!!!!!!!")
                        self.best_win_ratio = win_ratio
                        # update the best_policy
                        self.policy_value_net.save_model(
                            MODEL_DIR / self.name / 'train' / 'best'
                        )
                        self.policy_value_net.save_for_predict(
                            MODEL_DIR / self.name / 'production' / 'best'
                        )
                        if (self.best_win_ratio == 1.0 and
                                self.pure_mcts_playout_num < 5000):
                            self.pure_mcts_playout_num += 1000
                            self.best_win_ratio = 0.0
        except KeyboardInterrupt:
            logger.info('[bold red]quit.')

    @property
    def name_suffix(self):
        return f'{self.board_width}x{self.board_height}-{self.n_in_rows}'