"""Engine UCI don gian, dung model CNN (policy + value) train tu dataset random-vs-stockfish.
Chay: python uci_engine.py
Yeu cau: pip install torch python-chess
File chess_model.pt phai cung thu muc (hoac sua MODEL_PATH ben duoi)."""
import sys
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import chess
import numpy as np

MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "chess_model.pt")
ENGINE_NAME = "RandomVsSF-Net"
ENGINE_AUTHOR = "Ban tu train"

PIECE_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]


class ChessNet(nn.Module):
    def __init__(self, in_planes, channels, num_blocks):
        super().__init__()
        self.conv_in = nn.Conv2d(in_planes, channels, 3, padding=1)
        self.bn_in = nn.BatchNorm2d(channels)
        self.blocks = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(channels, channels, 3, padding=1),
                nn.BatchNorm2d(channels),
                nn.ReLU(),
                nn.Conv2d(channels, channels, 3, padding=1),
                nn.BatchNorm2d(channels),
            ) for _ in range(num_blocks)
        ])
        self.policy_conv = nn.Conv2d(channels, 32, 1)
        self.policy_fc = nn.Linear(32 * 8 * 8, 4096)
        self.value_conv = nn.Conv2d(channels, 8, 1)
        self.value_fc1 = nn.Linear(8 * 8 * 8, 64)
        self.value_fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = F.relu(self.bn_in(self.conv_in(x)))
        for block in self.blocks:
            x = F.relu(block(x) + x)
        p = F.relu(self.policy_conv(x)).flatten(1)
        policy_logits = self.policy_fc(p)
        v = F.relu(self.value_conv(x)).flatten(1)
        v = F.relu(self.value_fc1(v))
        value = torch.tanh(self.value_fc2(v))
        return policy_logits, value


def board_to_tensor(board, num_planes):
    planes = np.zeros((num_planes, 8, 8), dtype=np.float32)
    for color_idx, color in enumerate([chess.WHITE, chess.BLACK]):
        for piece_idx, piece_type in enumerate(PIECE_ORDER):
            for sq in board.pieces(piece_type, color):
                r, c = chess.square_rank(sq), chess.square_file(sq)
                planes[color_idx * 6 + piece_idx, r, c] = 1.0
    if board.turn == chess.WHITE:
        planes[12, :, :] = 1.0
    planes[13, :, :] = 1.0 if board.has_kingside_castling_rights(chess.WHITE) else 0.0
    planes[14, :, :] = 1.0 if board.has_queenside_castling_rights(chess.WHITE) else 0.0
    planes[15, :, :] = 1.0 if board.has_kingside_castling_rights(chess.BLACK) else 0.0
    planes[16, :, :] = 1.0 if board.has_queenside_castling_rights(chess.BLACK) else 0.0
    return planes


def select_move(model, board, device, num_planes, top_k=5):
    model.eval()
    x = torch.from_numpy(board_to_tensor(board, num_planes)).unsqueeze(0).to(device)
    with torch.no_grad():
        policy_logits, _ = model(x)
    policy_logits = policy_logits.squeeze(0).cpu().numpy()

    legal_moves = list(board.legal_moves)
    if not legal_moves:
        return None

    scored = [(mv, policy_logits[mv.from_square * 64 + mv.to_square]) for mv in legal_moves]
    scored.sort(key=lambda t: t[1], reverse=True)
    candidates = scored[:top_k]

    mover_is_white = board.turn == chess.WHITE
    best_move, best_value = None, -float("inf")
    for mv, _ in candidates:
        board.push(mv)
        x2 = torch.from_numpy(board_to_tensor(board, num_planes)).unsqueeze(0).to(device)
        with torch.no_grad():
            _, value = model(x2)
        v = value.item()
        v_for_mover = v if mover_is_white else -v
        board.pop()
        if v_for_mover > best_value:
            best_value = v_for_mover
            best_move = mv
    return best_move


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    ckpt = torch.load(MODEL_PATH, map_location=device)
    model = ChessNet(ckpt["in_planes"], ckpt["channels"], ckpt["num_blocks"]).to(device)
    model.load_state_dict(ckpt["model_state"])
    model.eval()
    num_planes = ckpt["in_planes"]

    board = chess.Board()

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        tokens = line.split()
        cmd = tokens[0]

        if cmd == "uci":
            print(f"id name {ENGINE_NAME}")
            print(f"id author {ENGINE_AUTHOR}")
            print("uciok")
            sys.stdout.flush()
        elif cmd == "isready":
            print("readyok")
            sys.stdout.flush()
        elif cmd in ("ucinewgame", "ucinewchess"):
            board = chess.Board()
        elif cmd == "position":
            idx = 1
            if tokens[idx] == "startpos":
                board = chess.Board()
                idx += 1
            elif tokens[idx] == "fen":
                fen = " ".join(tokens[idx + 1: idx + 7])
                board = chess.Board(fen)
                idx += 7
            if idx < len(tokens) and tokens[idx] == "moves":
                for mv_str in tokens[idx + 1:]:
                    board.push(chess.Move.from_uci(mv_str))
        elif cmd == "go":
            mv = select_move(model, board, device, num_planes)
            if mv is None:
                print("bestmove 0000")
            else:
                print(f"bestmove {mv.uci()}")
            sys.stdout.flush()
        elif cmd == "quit":
            break


if __name__ == "__main__":
    main()