// amazons_engine.h
extern "C" void batchLoadBoard(const int32_t* board_ptr, size_t rows, size_t cols);

// amazons_engine.cpp
#include <cstdint>
void batchLoadBoard(const int32_t* board_ptr, size_t rows, size_t cols) {
    // 假设 engine 提供一个快速接口直接从 int32_t* 填盘
    // 这里必须假设 board_ptr 指向 row*cols 连续内存
    for (size_t r = 0; r < rows; ++r)
        for (size_t c = 0; c < cols; ++c) {
            int32_t v = board_ptr[r*cols + c];
            // internalSetPiece(r, c, v);  // 内部超快接口
        }
}

import numpy as np
import amazons_engine

board = np.ascontiguousarray(board_numpy, dtype=np.int32)
# 强制 contiguous + int32（重要）
amazons_engine.batchLoadBoard(board)

from multiprocessing import get_context, shared_memory
import numpy as np
ctx = get_context('spawn')
task_queue = ctx.Queue()  # native queue
shm = shared_memory.SharedMemory(create=True, size=slot_count * slot_size)

# 每个 Actor 写入自己的 slot index
def actor_send(board_array, slot_idx):
    arr = np.ndarray((8,8), dtype=np.int32, buffer=shm.buf, offset=slot_idx*slot_bytes)
    np.copyto(arr, np.ascontiguousarray(board_array, dtype=np.int32))
    # 只 send slot_idx 和 metadata，避免拷贝
    task_queue.put({'slot': slot_idx, 'actor_id': my_id})

def worker_loop():
    while True:
        task = task_queue.get()  # blocking native queue
        slot_idx = task['slot']
        arr = np.ndarray((8,8), dtype=np.int32, buffer=shm.buf, offset=slot_idx*slot_bytes)
        # 直接把 arr 传给 batchLoadBoard 或者打包给 GPU
        amazons_engine.batchLoadBoard(arr)  # zero-copy-like
        # 回复 Actor：用另一个小队列发送 slot_idx/status

import psutil, os, time, random
p = psutil.Process(os.getpid())
# Linux: set to high priority (negative niceness)
p.nice(-10)

# Actor start jitter
time.sleep(random.uniform(0.1, 3.0))

# 每一步搜索开始前
mcts.clear_tree()

def board_to_onehot(board, me=1):
    # board: 8x8 ints, using {1: my, -1: opp, 3: arrow} or similar
    my_map = (board == me).astype(np.float32)
    opp_map = (board == -me).astype(np.float32)
    arrow_map = (board == 3).astype(np.float32)
    return np.stack([my_map, opp_map, arrow_map], axis=0)  # shape (3,8,8)

import torch
import torch.nn as nn
import torch.nn.functional as F
import sys
sys.path.append('..')

class ResBlock(nn.Module):
    """
    标准残差块：
    Structure: Conv3x3 -> BN -> ReLU -> Conv3x3 -> BN -> (+) -> ReLU
    """
    def __init__(self, num_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(num_channels)
        self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(num_channels)

    def forward(self, x):
        residual = x
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += residual
        return F.relu(out)

class AmazonsNNet(nn.Module):
    def __init__(self, game, args):
        super().__init__()

        self.board_x, self.board_y = game.getBoardSize()
        self.args = args
        self.flat_size = self.board_x * self.board_y

        
        # ResNet-40 对应 20 个残差块 (20 * 2 + 1 = 41 Layers)
        
        self.num_res_blocks = getattr(args, 'num_res_blocks', 20) 
        self.num_channels = args.num_channels # 建议 256

        # --- 1. 初始卷积层 ---
        # 输入: 3 通道 One-Hot (My, Op, Arrow)
        # 输出: num_channels
        self.conv1 = nn.Conv2d(3, self.num_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(self.num_channels)

        # --- 2. 残差塔 ---
        # 堆叠 20 个残差块
        self.res_layers = nn.ModuleList([ResBlock(self.num_channels) for _ in range(self.num_res_blocks)])

        # --- 3. 策略头 ---
        # 输入: (Batch, num_channels, 8, 8)
        self.fc_input_dim = self.num_channels * self.board_x * self.board_y
        
        self.fc1 = nn.Linear(self.fc_input_dim, 1024)
        self.fc2 = nn.Linear(1024, 512)

        # AlphaZero 标准是 3 个头分别输出 Logits
        self.fc_src = nn.Linear(512, self.flat_size)
        self.fc_dst = nn.Linear(512, self.flat_size)
        self.fc_arrow = nn.Linear(512, self.flat_size)

        # --- 4. 价值头---
        self.fc_val = nn.Linear(512, 1)

    def forward(self, s):
        # s: batch_size x 3 x 8 x 8
        
        # 1. 初始卷积
        s = F.relu(self.bn1(self.conv1(s)))

        # 2. 通过 ResNet-40 塔
        for layer in self.res_layers:
            s = layer(s)

        # 3. 展平
        s = s.view(-1, self.fc_input_dim)

        # 4. 共享全连接层
        s = F.dropout(F.relu(self.fc1(s)), p=self.args.dropout, training=self.training)
        s = F.dropout(F.relu(self.fc2(s)), p=self.args.dropout, training=self.training)

        # 5. 输出策略 
        pi_src = F.log_softmax(self.fc_src(s), dim=1)
        pi_dst = F.log_softmax(self.fc_dst(s), dim=1)
        pi_arrow = F.log_softmax(self.fc_arrow(s), dim=1)

        # 6. 输出价值 
        v = torch.tanh(self.fc_val(s))

        return pi_src, pi_dst, pi_arrow, v

torch.save({
    'epoch': epoch,
    'model_state_dict': net.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'rng_state': torch.get_rng_state(),
}, 'checkpoint.pth.tar')

# 加载
ckpt = torch.load('checkpoint.pth.tar', map_location=device)
net.load_state_dict(ckpt['model_state_dict'])
optimizer.load_state_dict(ckpt['optimizer_state_dict'])

# 在每一步 MCTS 结束后：
visit_counts = mcts.get_root_visit_counts()  # dict{action_idx: count}
total = sum(visit_counts.values())
pi_vector = [visit_counts.get(i, 0)/total for i in all_actions]  # soft labels

# decode pi_vector -> (src_map, dst_map, arr_map)
src_map, dst_map, arr_map = decode_pi_to_maps(pi_vector)

# 存入 trainExamples: (board_onehot, (src_map, dst_map, arr_map), value)
trainExamples.append((board_onehot, (src_map, dst_map, arr_map), value))

# 选择落子
if move_number < temp_threshold:
    action_idx = sample_from_distribution(pi_vector)  # tau=1
else:
    action_idx = argmax(pi_vector)  # tau=0