!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
"""
Chess AI — AlphaZero-style
Backend : PyTorch CUDA  (thay thế TinyGrad)
Checkpoint : pickle  (giữ nguyên)
Deps : torch, chess, numpy, tqdm
"""

EVAL_CACHE = {}
CACHE_MAX  = 100_000

import os, pickle, time, traceback
import chess, chess.pgn
import numpy as np
from tqdm.auto import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[Device] {DEVICE}  |  "
      f"{torch.cuda.get_device_name(0) if DEVICE.type == 'cuda' else 'CPU'}")

# =============================================================
# CONFIG
# =============================================================
CFG = {
    "mcts_sims":       16,
    "buffer_maxlen":   50_000,
    "min_buffer":      2_000,
    "batch_size":      512,
    "lr":              1e-4,
    "train_steps":     30,
    "save_every":      1,
    "pgn_limit":       2000,
    "res_blocks":      6,
    "channels":        256,
    "max_moves":       81,
    "checkpoint_path": "/content/drive/MyDrive/dataset stockfish/latest.pkl",
    "pgn_path":        "/content/drive/MyDrive/dataset stockfish/million_games.pgn",
    "log_path":        "/content/drive/MyDrive/dataset stockfish/training_log.txt",
}
CFG["pgn_chunk"]       = 500
CFG["pgn_offset_file"] = "/content/drive/MyDrive/dataset stockfish/pgn_offset.pkl"
TOTAL_STEPS_TARGET     = CFG["train_steps"]
os.makedirs(os.path.dirname(CFG["checkpoint_path"]), exist_ok=True)

# =============================================================
# ENCODING  (NumPy — CPU)
# =============================================================
MCTS_BATCH = 16

PIECE_IDX = {
    chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
    chess.ROOK: 3, chess.QUEEN:  4, chess.KING:   5,
}
PROMOTION_OFFSET = 4096
PROMOTION_MAP    = {chess.QUEEN: 0, chess.ROOK: 1,
                    chess.BISHOP: 2, chess.KNIGHT: 3}
ACTION_SIZE = 4672

ENCODE_CACHE = {}

def encode_board(board: chess.Board) -> np.ndarray:
    key = board.fen()
    if key in ENCODE_CACHE:
        return ENCODE_CACHE[key]
    x = np.zeros((12, 8, 8), dtype=np.float32)
    for sq, piece in board.piece_map().items():
        row, col = divmod(sq, 8)
        offset = 0 if piece.color == chess.WHITE else 6
        x[PIECE_IDX[piece.piece_type] + offset][row][col] = 1.0
    ENCODE_CACHE[key] = x
    return x

def move_to_idx(move: chess.Move) -> int:
    if move.promotion is not None:
        return PROMOTION_OFFSET + move.to_square * 4 + PROMOTION_MAP[move.promotion]
    return move.from_square * 64 + move.to_square

# =============================================================
# MODEL  (PyTorch)
# =============================================================
class ResBlock(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.c1  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.c2  = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = F.relu(self.bn1(self.c1(x)))
        x = self.bn2(self.c2(x))
        return F.relu(x + residual)


class ChessNet(nn.Module):
    def __init__(self, channels: int = 256, num_blocks: int = 6):
        super().__init__()
        self.stem_conv = nn.Conv2d(12, channels, 3, padding=1, bias=False)
        self.stem_bn   = nn.BatchNorm2d(channels)
        self.blocks    = nn.ModuleList([ResBlock(channels) for _ in range(num_blocks)])

        # Policy head
        self.policy_conv = nn.Conv2d(channels, 32, 1, bias=False)
        self.policy_bn   = nn.BatchNorm2d(32)
        self.policy_fc   = nn.Linear(32 * 8 * 8, ACTION_SIZE)

        # Value head
        self.value_conv  = nn.Conv2d(channels, 8, 1, bias=False)
        self.value_bn    = nn.BatchNorm2d(8)
        self.value_fc1   = nn.Linear(8 * 8 * 8, 256)
        self.value_fc2   = nn.Linear(256, 1)

    def forward(self, x: torch.Tensor):
        x = F.relu(self.stem_bn(self.stem_conv(x)))
        for block in self.blocks:
            x = block(x)

        p = F.relu(self.policy_bn(self.policy_conv(x)))
        p = self.policy_fc(p.flatten(1))

        v = F.relu(self.value_bn(self.value_conv(x)))
        v = F.relu(self.value_fc1(v.flatten(1)))
        v = torch.tanh(self.value_fc2(v))
        return p, v


model     = ChessNet(CFG["channels"], CFG["res_blocks"]).to(DEVICE)
optimizer = Adam(model.parameters(), lr=CFG["lr"])

# =============================================================
# REPLAY BUFFER
# =============================================================
class ReplayBuffer:
    def __init__(self, maxlen: int):
        self.maxlen   = maxlen
        self.states   = np.zeros((maxlen, 12, 8, 8), dtype=np.float32)
        self.policies = np.zeros((maxlen, ACTION_SIZE), dtype=np.float32)
        self.values   = np.zeros((maxlen, 1), dtype=np.float32)
        self.ptr  = 0
        self.size = 0

    def extend(self, samples):
        for state, pi, z in samples:
            self.states[self.ptr]    = state
            self.policies[self.ptr]  = pi
            self.values[self.ptr, 0] = z
            self.ptr  = (self.ptr + 1) % self.maxlen
            self.size = min(self.size + 1, self.maxlen)

    def sample_batch(self, batch_size: int):
        idx = np.random.randint(0, self.size, batch_size)
        return self.states[idx], self.policies[idx], self.values[idx]

    def __len__(self):
        return self.size

# =============================================================
# CHECKPOINT  (pickle — giữ nguyên format)
# =============================================================
def save_checkpoint(games_played: int):
    data = {
        "games":       games_played,
        "total_steps": total_steps,
        "weights":     {k: v.cpu() for k, v in model.state_dict().items()},
        "optimizer":   optimizer.state_dict(),
    }
    with open(CFG["checkpoint_path"], "wb") as f:
        pickle.dump(data, f)
    tqdm.write(f"  [CKPT] saved → games={games_played}, steps={total_steps:,}")


def load_checkpoint():
    path = CFG["checkpoint_path"]
    if not os.path.exists(path):
        tqdm.write("  [CKPT] không tìm thấy checkpoint, bắt đầu từ đầu.")
        return 0, 0
    with open(path, "rb") as f:
        data = pickle.load(f)
    model.load_state_dict(data["weights"])
    optimizer.load_state_dict(data["optimizer"])
    games = data.get("games", 0)
    steps = data.get("total_steps", 0)
    tqdm.write(f"  [CKPT] loaded → games={games}, steps={steps:,}")
    return games, steps

# =============================================================
# PGN LOADING
# =============================================================
def load_pgn_offset():
    if os.path.exists(CFG["pgn_offset_file"]):
        with open(CFG["pgn_offset_file"], "rb") as f:
            return pickle.load(f)
    return 0


def load_pgn_chunk(buf: ReplayBuffer, path: str, max_games: int = 500):
    if not os.path.exists(path):
        return 0, None
    offset = load_pgn_offset()
    count  = 0
    with open(path, "r") as f:
        f.seek(offset)
        while count < max_games:
            game = chess.pgn.read_game(f)
            if game is None:
                break
            result  = game.headers.get("Result", "1/2-1/2")
            white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
            board   = game.board()
            samples = []
            for move in game.mainline_moves():
                pi = np.zeros(ACTION_SIZE, dtype=np.float32)
                pi[move_to_idx(move)] = 1.0
                z = white_z if board.turn == chess.WHITE else -white_z
                samples.append((encode_board(board), pi, np.float32(z)))
                board.push(move)
            buf.extend(samples)
            count += 1
        new_offset = f.tell()
    with open(CFG["pgn_offset_file"], "wb") as f:
        pickle.dump(new_offset, f)
    tqdm.write(f"[PGN] loaded chunk={count}, buffer={len(buf):,}")
    return count, new_offset


def load_pgn(buf: ReplayBuffer, path: str, limit: int = 2000):
    if not os.path.exists(path):
        tqdm.write(f"  [PGN] không tìm thấy {path}, bỏ qua.")
        return
    count = 0
    with open(path) as f:
        pbar = tqdm(total=limit, desc="Loading PGN")
        while count < limit:
            game = chess.pgn.read_game(f)
            if game is None:
                break
            result  = game.headers.get("Result", "1/2-1/2")
            white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
            board   = game.board()
            samples = []
            for move in game.mainline_moves():
                pi = np.zeros(ACTION_SIZE, dtype=np.float32)
                pi[move_to_idx(move)] = 1.0
                z = white_z if board.turn == chess.WHITE else -white_z
                samples.append((encode_board(board), pi, np.float32(z)))
                board.push(move)
            buf.extend(samples)
            count += 1
            pbar.update(1)
        pbar.close()
    tqdm.write(f"  [PGN] loaded {count} games, buffer={len(buf):,}")

# =============================================================
# INFERENCE  (torch.no_grad)
# =============================================================
def _softmax_masked(logits_np: np.ndarray, mask: np.ndarray) -> np.ndarray:
    p = logits_np.astype(np.float64)
    p = np.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0)
    p -= p.max()
    np.exp(p, out=p)
    p *= mask
    total = p.sum()
    if total > 0:
        p /= total
    else:
        p = mask / mask.sum()
    return p


def predict(board: chess.Board):
    fen    = board.fen()
    cached = EVAL_CACHE.get(fen)
    if cached is not None:
        return cached

    x = torch.from_numpy(encode_board(board)[None]).to(DEVICE)  # (1,12,8,8)
    model.eval()
    with torch.no_grad():
        logits, v = model(x)

    p_np     = logits[0].cpu().numpy()
    v_scalar = float(v[0, 0].cpu().item())

    mask = np.zeros(ACTION_SIZE, dtype=np.float64)
    for m in board.legal_moves:
        mask[move_to_idx(m)] = 1.0

    p = _softmax_masked(p_np, mask)

    result = (p, v_scalar)
    if len(EVAL_CACHE) >= CACHE_MAX:
        EVAL_CACHE.clear()
    EVAL_CACHE[fen] = result
    return result


def predict_batch(boards):
    n      = len(boards)
    states = np.empty((n, 12, 8, 8), dtype=np.float32)
    for i, b in enumerate(boards):
        states[i] = encode_board(b)

    x = torch.from_numpy(states).to(DEVICE)
    model.eval()
    with torch.no_grad():
        logits, values = model(x)

    logits_np = logits.cpu().numpy().astype(np.float64, copy=False)
    values_np = values.cpu().numpy()

    outputs = []
    for i, board in enumerate(boards):
        mask = np.zeros(ACTION_SIZE, dtype=np.float64)
        for m in board.legal_moves:
            mask[move_to_idx(m)] = 1.0
        p = _softmax_masked(logits_np[i], mask)
        outputs.append((p.copy(), float(values_np[i, 0])))
    return outputs

# =============================================================
# MCTS  (CPU — giữ nguyên)
# =============================================================
class MCTSNode:
    __slots__ = ("board", "prior", "parent", "children", "visits", "total_value")

    def __init__(self, board: chess.Board, prior: float = 0.0, parent=None):
        self.board       = board
        self.prior       = prior
        self.parent      = parent
        self.children    = {}
        self.visits      = 0
        self.total_value = 0.0

    @property
    def q(self) -> float:
        return self.total_value / self.visits if self.visits > 0 else 0.0

    def uct_score(self) -> float:
        c_puct        = 1.5
        parent_visits = self.parent.visits if self.parent else 1
        u = c_puct * self.prior * (parent_visits ** 0.5) / (1 + self.visits)
        return self.q + u


def terminal_value(board: chess.Board) -> float:
    outcome = board.outcome()
    if outcome is None or outcome.winner is None:
        return 0.0
    return -1.0


def expand_node(node: MCTSNode, policy: np.ndarray):
    for move in node.board.legal_moves:
        if move not in node.children:
            child_board = node.board.copy(stack=False)
            child_board.push(move)
            node.children[move] = MCTSNode(
                child_board, prior=float(policy[move_to_idx(move)]), parent=node)


def backup_path(path, value: float):
    for n in reversed(path):
        n.visits      += 1
        n.total_value += value
        value = -value


def mcts_search(root: MCTSNode, num_sims: int = 64):
    remaining = num_sims
    while remaining > 0:
        batch_size = min(MCTS_BATCH, remaining)
        leaves, paths = [], []

        for _ in range(batch_size):
            node = root
            path = [node]
            while node.children and not node.board.is_game_over():
                node = max(node.children.values(), key=lambda n: n.uct_score())
                path.append(node)
            leaves.append(node)
            paths.append(path)

        boards_to_eval, eval_indices = [], []
        values   = [0.0] * batch_size
        policies = [None] * batch_size

        for i, node in enumerate(leaves):
            if node.board.is_game_over():
                values[i] = terminal_value(node.board)
            else:
                boards_to_eval.append(node.board)
                eval_indices.append(i)

        if boards_to_eval:
            outputs = predict_batch(boards_to_eval)
            for idx, (policy, value) in zip(eval_indices, outputs):
                policies[idx] = policy
                values[idx]   = value

        for i, node in enumerate(leaves):
            if not node.board.is_game_over():
                policy = policies[i]
                if node == root:
                    legal = list(node.board.legal_moves)
                    if legal:
                        noise = np.random.dirichlet([0.3] * len(legal))
                        for j, move in enumerate(legal):
                            k = move_to_idx(move)
                            policy[k] = 0.75 * policy[k] + 0.25 * noise[j]
                expand_node(node, policy)
            backup_path(paths[i], values[i])

        remaining -= batch_size

# =============================================================
# SELF-PLAY
# =============================================================
def play_one_game(temperature: float = 1.0, game_id: int = 0):
    board    = chess.Board()
    history  = []
    move_num = 0

    move_bar = tqdm(
        total=CFG["max_moves"], desc=f"  ♟ Game #{game_id}",
        leave=False, dynamic_ncols=True)

    try:
        while not board.is_game_over() and move_num < CFG["max_moves"]:
            root = MCTSNode(board.copy(stack=False))
            mcts_search(root, CFG["mcts_sims"])

            moves_list = list(root.children.keys())
            if not moves_list:
                break

            visits = np.array(
                [root.children[m].visits for m in moves_list], dtype=np.float32)
            s = visits.sum()
            if s <= 0 or np.isnan(s):
                visits[:] = 1.0

            idxs = np.array([move_to_idx(m) for m in moves_list], dtype=np.int32)
            pi   = np.zeros(ACTION_SIZE, dtype=np.float32)
            pi[idxs] = visits
            s = pi.sum()
            if s > 0 and not np.isnan(s):
                pi /= s

            history.append((encode_board(board), pi, board.turn))

            if move_num < 30 and temperature > 0:
                probs = visits ** (1.0 / temperature)
                s = probs.sum()
                if s <= 0 or np.isnan(s):
                    probs = np.ones(len(visits), dtype=np.float32)
                probs /= probs.sum()
                move = moves_list[np.random.choice(len(moves_list), p=probs)]
            else:
                move = moves_list[int(np.argmax(visits))]

            board.push(move)
            move_num += 1
            move_bar.update(1)
            move_bar.set_postfix(
                turn="W" if board.turn == chess.BLACK else "B",
                move=move.uci(), sims=CFG["mcts_sims"])
    finally:
        move_bar.close()

    result  = board.result() if board.is_game_over() else "1/2-1/2"
    white_z = 1.0 if result == "1-0" else -1.0 if result == "0-1" else 0.0
    samples = [
        (state, pi, np.float32(white_z if color == chess.WHITE else -white_z))
        for state, pi, color in history
    ]
    return samples, result

# =============================================================
# TRAINING STEP
# =============================================================
def train_step(buf: ReplayBuffer):
    if len(buf) < CFG["min_buffer"]:
        return None

    model.train()
    train_bar = tqdm(
        total=CFG["train_steps"], desc="  🔧 Training",
        leave=False, dynamic_ncols=True)

    loss = loss_policy = loss_value = 0.0

    for _ in range(CFG["train_steps"]):
        s, p, z = buf.sample_batch(CFG["batch_size"])

        S = torch.from_numpy(np.ascontiguousarray(s)).to(DEVICE)   # (B,12,8,8)
        P = torch.from_numpy(np.ascontiguousarray(p)).to(DEVICE)   # (B,ACTION_SIZE)
        Z = torch.from_numpy(np.ascontiguousarray(z)).to(DEVICE)   # (B,1)

        logits, v = model(S)

        # Policy loss: stable cross-entropy via log_softmax
        log_prob   = F.log_softmax(logits, dim=1)
        loss_policy = -(P * log_prob).sum(dim=1).mean()

        # Value loss: MSE
        loss_value = F.mse_loss(v, Z)

        loss = loss_policy + loss_value

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_bar.update(1)
        train_bar.set_postfix(loss=f"{loss.item():.4f}")

    train_bar.close()

    return {
        "loss":        loss.item(),
        "loss_policy": loss_policy.item(),
        "loss_value":  loss_value.item(),
        "steps":       CFG["train_steps"],
    }

# =============================================================
# MAIN
# =============================================================
if __name__ == "__main__":
    step_since_pgn = 0

    print("=" * 60)
    print("  Chess AI — AlphaZero-style  |  PyTorch CUDA")
    print(f"  Device : {DEVICE}")
    print("=" * 60)

    buf = ReplayBuffer(CFG["buffer_maxlen"])
    games_played, total_steps = load_checkpoint()

    if games_played == 0:
        load_pgn(buf, CFG["pgn_path"], CFG["pgn_limit"])
    load_pgn_chunk(buf, CFG["pgn_path"], max_games=500)

    print(f"\n  Bắt đầu training từ game #{games_played + 1}")
    print(f"  Buffer: {len(buf):,} | Min train: {CFG['min_buffer']:,}\n")

    game_bar = tqdm(initial=games_played, unit="game",
                    desc="🎮 Games", dynamic_ncols=True)
    step_bar = tqdm(
        total=TOTAL_STEPS_TARGET, initial=total_steps,
        unit="step", desc="Steps", dynamic_ncols=True,
        bar_format="{desc}: {n_fmt}/{total_fmt} |{bar}| {elapsed}<{remaining}")

    while True:
        t0 = time.time()
        step_since_pgn += CFG["train_steps"]
        if step_since_pgn >= CFG["pgn_chunk"]:
            load_pgn_chunk(buf, CFG["pgn_path"], CFG["pgn_chunk"])
            step_since_pgn = 0

        # ── Self-play ────────────────────────────────────────
        try:
            samples, result = play_one_game(temperature=1.0, game_id=games_played + 1)
        except Exception as e:
            ts  = time.strftime("%Y-%m-%d %H:%M:%S")
            msg = f"[ERROR] self-play: {e}\n{traceback.format_exc()}"
            tqdm.write(f"[{ts}] {msg}")
            try:
                with open(CFG["log_path"], "a") as f:
                    f.write(f"[{ts}] {msg}\n")
            except Exception:
                pass
            continue

        buf.extend(samples)
        games_played += 1
        elapsed_game  = time.time() - t0

        # ── Train ────────────────────────────────────────────
        t1    = time.time()
        info  = train_step(buf)
        elapsed_train = time.time() - t1

        if info:
            step_bar.update(1)
            total_steps += 1
            step_bar.set_postfix(
                progress=f"{total_steps}/{TOTAL_STEPS_TARGET}",
                loss=f"{info['loss']:.4f}",
                p=f"{info['loss_policy']:.4f}",
                v=f"{info['loss_value']:.4f}",
                t=f"{elapsed_train:.1f}s",
            )

        game_bar.update(1)
        game_bar.set_postfix(
            result=result, moves=len(samples),
            buf=f"{len(buf):,}",
            loss=f"{info['loss']:.4f}" if info else "warmup",
            t=f"{elapsed_game:.1f}s",
        )

        # ── Log ──────────────────────────────────────────────
        ts   = time.strftime("%Y-%m-%d %H:%M:%S")
        line = (
            f"game={games_played:5d} | {result} | moves={len(samples):3d} | "
            f"buf={len(buf):,} | steps={total_steps:,} | "
            + (f"loss={info['loss']:.4f} "
               f"(p={info['loss_policy']:.4f} v={info['loss_value']:.4f})"
               if info else "warmup")
            + f" | game={elapsed_game:.1f}s train={elapsed_train:.1f}s"
        )
        try:
            with open(CFG["log_path"], "a") as f:
                f.write(f"[{ts}] {line}\n")
        except Exception:
            pass

        if games_played % CFG["save_every"] == 0:
            save_checkpoint(games_played)