app.py

import os
import math
import random
import time
from typing import List, Tuple, Dict, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import chess
import chess.svg
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# Check for GPU
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Precomputed sq → (row, col)
_SQ_ROW   = np.array([sq // 8 for sq in range(64)], dtype=np.int32)
_SQ_COL   = np.array([sq %  8 for sq in range(64)], dtype=np.int32)
_SQ_ROW_B = (7 - _SQ_ROW)
_SQ_COL_B = (7 - _SQ_COL)


# ── Config ────────────────────────────────────────────────────────────────────────
class ChessRLConfig:
    board_channels: int = 18
    max_moves:      int = 20480
    channels:       int = 32      
    num_res_blocks: int = 4       
    policy_planes:  int = 2       
    mcts_c_puct:      float = 1.5
    dirichlet_alpha:  float = 0.3
    dirichlet_epsilon: float = 0.25


# ── Move ↔ index ──────────────────────────────────────────────────────────────────
class MoveIndexConverter:
    def __init__(self):
        self.promo_map    = {None: 0, chess.QUEEN: 1, chess.ROOK: 2,
                             chess.BISHOP: 3, chess.KNIGHT: 4}
        self.idx_to_promo = {v: k for k, v in self.promo_map.items()}

    def move_to_code(self, move: chess.Move) -> int:
        return move.from_square * 320 + move.to_square * 5 + self.promo_map[move.promotion]

    def code_to_move(self, code: int) -> chess.Move:
        from_sq = code // 320
        rest    = code % 320
        return chess.Move(from_sq, rest // 5, promotion=self.idx_to_promo[rest % 5])


# ── Board encoding ───────────────────────────────────────────────────────────────
def encode_board(board: chess.Board, config: ChessRLConfig,
                 perspective: Optional[chess.Color] = None) -> torch.Tensor:
    if perspective is None:
        perspective = board.turn

    x       = np.zeros((config.board_channels, 8, 8), dtype=np.float32)
    row_lut = _SQ_ROW_B if perspective == chess.BLACK else _SQ_ROW
    col_lut = _SQ_COL_B if perspective == chess.BLACK else _SQ_COL

    for sq, piece in board.piece_map().items():
        offset = 0 if piece.color == perspective else 6
        x[offset + piece.piece_type - 1, row_lut[sq], col_lut[sq]] = 1.0

    x[12] = 1.0
    w, b = chess.WHITE, chess.BLACK
    if perspective == w:
        x[13] = float(board.has_kingside_castling_rights(w))
        x[14] = float(board.has_queenside_castling_rights(w))
        x[15] = float(board.has_kingside_castling_rights(b))
        x[16] = float(board.has_queenside_castling_rights(b))
    else:
        x[13] = float(board.has_kingside_castling_rights(b))
        x[14] = float(board.has_queenside_castling_rights(b))
        x[15] = float(board.has_kingside_castling_rights(w))
        x[16] = float(board.has_queenside_castling_rights(w))
    x[17] = board.halfmove_clock / 100.0

    return torch.from_numpy(x).contiguous()


# ── Tiny Model Architecture ──────────────────────────────────────────────────────
class ResidualBlock(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1   = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2   = nn.BatchNorm2d(channels)

    def forward(self, x):
        r = x
        x = F.relu(self.bn1(self.conv1(x)), inplace=True)
        x = self.bn2(self.conv2(x))
        return F.relu(x + r, inplace=True)


class ChessPolicyValueNet(nn.Module):
    def __init__(self, config: ChessRLConfig):
        super().__init__()
        C_in, C, P = config.board_channels, config.channels, config.policy_planes
        self.conv_in    = nn.Conv2d(C_in, C, 3, padding=1, bias=False)
        self.bn_in      = nn.BatchNorm2d(C)
        self.res_blocks = nn.ModuleList(
            [ResidualBlock(C) for _ in range(config.num_res_blocks)]
        )
        self.conv_p = nn.Conv2d(C, P, 1, bias=False)
        self.bn_p   = nn.BatchNorm2d(P)
        self.fc_p   = nn.Linear(P * 64, config.max_moves)

        self.conv_v = nn.Conv2d(C, 4, 1, bias=False)
        self.bn_v   = nn.BatchNorm2d(4)
        self.fc_v1  = nn.Linear(4 * 64, 64)
        self.fc_v2  = nn.Linear(64, 1)

    def forward(self, x):
        x = F.relu(self.bn_in(self.conv_in(x)), inplace=True)
        for blk in self.res_blocks:
            x = blk(x)

        p = F.relu(self.bn_p(self.conv_p(x)), inplace=True).flatten(1)
        p = self.fc_p(p)

        v = F.relu(self.bn_v(self.conv_v(x)), inplace=True).flatten(1)
        v = F.relu(self.fc_v1(v), inplace=True)
        v = torch.tanh(self.fc_v2(v))
        return p, v.squeeze(-1)


# ── MCTS Implementation ──────────────────────────────────────────────────────────
class MCTSNode:
    __slots__ = ('parent', 'child_idx_in_parent', 'move_from_parent', '_board',
                 'move_codes', 'priors', 'child_N', 'child_W', 'children',
                 'is_expanded', '_game_over')

    def __init__(self, board, parent=None, move_from_parent=None, child_idx=-1):
        self.parent               = parent
        self.child_idx_in_parent  = child_idx
        self.move_from_parent     = move_from_parent
        self._board               = board
        self.move_codes  = None   
        self.priors      = None   
        self.child_N     = None   
        self.child_W     = None   
        self.children    = None   
        self.is_expanded = False
        self._game_over  = None

    @property
    def board(self) -> chess.Board:
        if self._board is None:
            self._board = self.parent.board.copy()
            self._board.push(self.move_from_parent)
        return self._board

    def expand(self, move_codes: list, priors: list):
        n               = len(move_codes)
        self.move_codes = np.array(move_codes, dtype=np.int32)
        self.priors     = np.array(priors,     dtype=np.float32)
        self.child_N    = np.zeros(n, dtype=np.float32)
        self.child_W    = np.zeros(n, dtype=np.float32)
        self.children   = [None] * n
        self.is_expanded = True

    def best_child_idx(self, c_puct: float) -> int:
        sqrt_N = math.sqrt(float(self.child_N.sum()) + 1e-8)
        Q      = self.child_W / (self.child_N + 1e-8)
        U      = c_puct * self.priors * sqrt_N / (1.0 + self.child_N)
        return int(np.argmax(Q + U))

    def is_leaf(self) -> bool:
        return not self.is_expanded

    def game_over(self) -> bool:
        if self._game_over is None:
            self._game_over = self.board.is_game_over()
        return self._game_over


def backup_path(path: list, value: float):
    v = value
    for node, child_idx in reversed(path):
        v = -v
        node.child_N[child_idx] += 1
        node.child_W[child_idx] += v


def _terminal_value(board: chess.Board) -> float:
    res = board.result()
    if res == "1-0": return 1.0 if board.turn == chess.BLACK else -1.0
    if res == "0-1": return 1.0 if board.turn == chess.WHITE else -1.0
    return 0.0


# ── Model Downloader & Initialization ─────────────────────────────────────────────
print("Downloading weights from HuggingFace...")
try:
    weights_path = hf_hub_download(
        repo_id="FlameF0X/ChessRLM-tiny", 
        filename="model.safetensors"
    )
    print(f"Model successfully downloaded: {weights_path}")
except Exception as e:
    print(f"Error downloading weights: {e}. Running with initialized weights as fallback.")
    weights_path = None

config = ChessRLConfig()
net = ChessPolicyValueNet(config).to(DEVICE)

if weights_path:
    raw_state_dict = load_file(weights_path)
    
    # Strip any "module." prefix if present in the state_dict (due to DataParallel training)
    clean_state_dict = {}
    for k, v in raw_state_dict.items():
        if k.startswith("module."):
            clean_state_dict[k[7:]] = v
        else:
            clean_state_dict[k] = v
            
    net.load_state_dict(clean_state_dict)
    print("Model loaded successfully after fixing the key names!")
else:
    print("Loaded with random weights.")

net.eval()
move_converter = MoveIndexConverter()


# ── MCTS Single Decision Engine ──────────────────────────────────────────────────
def get_ai_move(
    board: chess.Board, 
    simulations: int, 
    temperature: float
) -> Tuple[chess.Move, float, List[Tuple[str, float]]]:
    """Runs MCTS search to find the best move, returning stats and evaluations."""
    root = MCTSNode(board.copy())
    
    for sim in range(simulations):
        node = root
        path = []
        
        # 1. Selection
        while not node.is_leaf() and not node.game_over():
            ci    = node.best_child_idx(config.mcts_c_puct)
            child = node.children[ci]
            if child is None:
                move  = move_converter.code_to_move(int(node.move_codes[ci]))
                child = MCTSNode(None, parent=node, move_from_parent=move, child_idx=ci)
                node.children[ci] = child
            path.append((node, ci))
            node = child

        # 2. Expansion / Evaluation
        if node.game_over():
            backup_path(path, _terminal_value(node.board))
        else:
            # Neural network prediction
            state_tensor = encode_board(node.board, config, node.board.turn).unsqueeze(0).to(DEVICE)
            with torch.no_grad():
                logits, value_tensor = net(state_tensor)
            
            policies = F.softmax(logits.float(), dim=-1).cpu().numpy()[0]
            val = float(value_tensor.float().cpu().numpy()[0])
            
            legal = list(node.board.legal_moves)
            if not legal:
                backup_path(path, 0.0)
                continue

            codes, priors_raw = [], []
            total_p = 0.0
            for move in legal:
                code = move_converter.move_to_code(move)
                p    = float(policies[code])
                codes.append(code)
                priors_raw.append(p)
                total_p += p

            if total_p > 1e-8:
                priors_norm = [p / total_p for p in priors_raw]
            else:
                priors_norm = [1.0 / len(codes)] * len(codes)

            # Dirichlet noise at root
            if node is root and sim == 0:
                noise = np.random.dirichlet([config.dirichlet_alpha] * len(codes))
                eps = config.dirichlet_epsilon
                priors_norm = [(1 - eps) * p + eps * float(n) for p, n in zip(priors_norm, noise)]

            node.expand(codes, priors_norm)
            backup_path(path, val)

    # Calculate Move Probabilities based on visit counts
    if root.move_codes is None or root.child_N.sum() == 0:
        fallback_move = random.choice(list(board.legal_moves))
        return fallback_move, 0.0, [("Random Fallback", 1.0)]

    total_N = float(root.child_N.sum()) + 1e-8
    visit_probs = root.child_N / total_N
    
    # Selection based on temperature
    if temperature > 1e-6:
        inv_t = 1.0 / temperature
        scaled = np.power(visit_probs, inv_t)
        scaled_sum = np.sum(scaled)
        if scaled_sum > 1e-8:
            choice_idx = np.random.choice(len(root.move_codes), p=scaled / scaled_sum)
        else:
            choice_idx = np.argmax(visit_probs)
    else:
        choice_idx = np.argmax(visit_probs)

    chosen_move_code = int(root.move_codes[choice_idx])
    best_move = move_converter.code_to_move(chosen_move_code)

    # Predicted value (win evaluation) from MCTS point-of-view
    win_rate_est = float(root.child_W[choice_idx] / (root.child_N[choice_idx] + 1e-8))

    # Top move evaluation lists
    top_moves = []
    sorted_indices = np.argsort(visit_probs)[::-1][:3]
    for idx in sorted_indices:
        mv = move_converter.code_to_move(int(root.move_codes[idx]))
        top_moves.append((board.san(mv), float(visit_probs[idx])))

    return best_move, win_rate_est, top_moves


# ── State Management ──────────────────────────────────────────────────────────────
class GameSession:
    def __init__(self):
        self.board = chess.Board()
        self.history = []
        self.player_color = chess.WHITE  # White by default
        self.game_log = []

    def reset(self, play_as_black: bool = False):
        self.board = chess.Board()
        self.history = []
        self.player_color = chess.BLACK if play_as_black else chess.WHITE
        self.game_log = [f"Game initialized. You are {'Black' if play_as_black else 'White'}."]

    def make_move(self, move_san: str):
        if self.board.is_game_over():
            return
        try:
            move = self.board.parse_san(move_san)
            self.history.append(self.board.copy())
            self.board.push(move)
            self.game_log.append(f"You: {move_san}")
        except ValueError:
            pass

    def make_ai_move(self, simulations: int, temp: float):
        if self.board.is_game_over():
            return "Game Over", []
        
        move, val, top_moves = get_ai_move(self.board, simulations, temp)
        san = self.board.san(move)
        self.history.append(self.board.copy())
        self.board.push(move)
        
        player_side = "White AI" if self.board.turn == chess.BLACK else "Black AI"
        self.game_log.append(f"{player_side}: {san}")
        
        # Format evaluation metrics
        color_eval = "Advantage: AI" if val > 0.05 else ("Advantage: Opponent" if val < -0.05 else "Equal")
        eval_str = f"Estimated Score: {val:+.2f} ({color_eval})"
        
        top_str = " | ".join([f"{m}: {p:.1%}" for m, p in top_moves])
        return eval_str, top_moves

    def undo(self):
        if self.history:
            self.board = self.history.pop()
            if self.game_log:
                self.game_log.pop()
            # If AI also moved, undo it as well to restore to player's turn
            if self.history and len(self.history) % 2 != 0:
                self.board = self.history.pop()
                if self.game_log:
                    self.game_log.pop()


# Persistent Session
session = GameSession()


# ── UI Rendering & Interactions ────────────────────────────────────────────────────
def get_svg_board(selected_sq=None) -> str:
    # Highlight last move if available
    last_move = session.board.peek() if session.board.move_stack else None
    check_sq = session.board.king(session.board.turn) if session.board.is_check() else None
    
    # Flip board if playing as Black
    flipped = (session.player_color == chess.BLACK)
    
    svg = chess.svg.board(
        board=session.board,
        lastmove=last_move,
        check=check_sq,
        flipped=flipped,
        size=450,
        colors={
            'square light': '#f0d9b5',
            'square dark': '#b58863',
            'square light lastmove': '#cdd26a',
            'square dark lastmove': '#aaa23a',
            'margin': '#1e293b'
        }
    )
    return f"<div style='display:flex; justify-content:center;'>{svg}</div>"


def get_game_status() -> str:
    if session.board.is_game_over():
        res = session.board.result()
        if res == "1-0": return "### 🏆 Game Over: White wins!"
        if res == "0-1": return "### 🏆 Game Over: Black wins!"
        return "### 🤝 Game Over: Draw!"
    
    turn = "White" if session.board.turn == chess.WHITE else "Black"
    is_check = " (In Check!)" if session.board.is_check() else ""
    return f"### Turn: **{turn}**{is_check}"


def update_ui(ai_eval="", top_moves_list=[]):
    # Prepare legal moves in SAN format for easy selection
    legal_sans = [session.board.san(m) for m in session.board.legal_moves]
    legal_sans.sort()
    
    # Render logs in reverse order (newest on top)
    logs_rendered = "\n".join([f"- {l}" for l in reversed(session.game_log)])
    
    analysis_md = ""
    if ai_eval:
        analysis_md += f"💡 **Evaluation**: {ai_eval}\n\n"
    if top_moves_list:
        analysis_md += "📈 **Top Considered Moves (Priorities)**:\n"
        for m, p in top_moves_list:
            analysis_md += f"- **{m}**: {p:.1%}\n"
            
    return (
        get_svg_board(),
        get_game_status(),
        gr.Dropdown(choices=legal_sans, value=None, label="Select Your Move", interactive=not session.board.is_game_over()),
        logs_rendered,
        analysis_md
    )


# ── Gradio Handlers ──────────────────────────────────────────────────────────────
def make_player_move(move_san, simulations, temp):
    if not move_san:
        return update_ui()
    
    # Play Player move
    session.make_move(move_san)
    
    # Check if game is over before AI moves
    if session.board.is_game_over():
        return update_ui()
        
    # Play AI move
    eval_str, top_moves = session.make_ai_move(int(simulations), float(temp))
    return update_ui(eval_str, top_moves)


def handle_reset(play_as_black, simulations, temp):
    session.reset(play_as_black=play_as_black)
    eval_str, top_moves = "", []
    
    # If starting as Black, AI goes first immediately
    if play_as_black:
        eval_str, top_moves = session.make_ai_move(int(simulations), float(temp))
        
    return update_ui(eval_str, top_moves)


def handle_undo():
    session.undo()
    return update_ui()


def handle_ai_vs_ai(simulations, temp):
    if session.board.is_game_over():
        return update_ui()
    eval_str, top_moves = session.make_ai_move(int(simulations), float(temp))
    return update_ui(eval_str, top_moves)


# ── CSS Styling ───────────────────────────────────────────────────────────────────
custom_css = """
body { background-color: #0f172a; color: #f8fafc; }
.gradio-container { max-width: 1100px !important; margin: auto; }
footer { display: none !important; }
"""


# ── Gradio Layout ─────────────────────────────────────────────────────────────────
# Note: theme and css are passed in launch() for Gradio 6.0+ compatibility
with gr.Blocks(title="ChessRLM-Tiny Space") as demo:
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1 style="color: #38bdf8; font-size: 2.5rem; font-weight: 800; margin-bottom: 5px;">♟️ ChessRLM-Tiny</h1>
        <p style="color: #94a3b8; font-size: 1.1rem;">Play against a ultra-light (~2.7M parameter) AlphaZero RL agent trained in PyTorch.</p>
    </div>
    """)
    
    with gr.Row():
        # Left Panel - Board Display
        with gr.Column(scale=5):
            board_html = gr.HTML(get_svg_board())
            status_markdown = gr.Markdown(get_game_status())
            
        # Right Panel - Gameplay & Settings
        with gr.Column(scale=4):
            with gr.Tab("Game Controls"):
                move_dropdown = gr.Dropdown(
                    choices=sorted([session.board.san(m) for m in session.board.legal_moves]),
                    label="Select Your Move",
                    interactive=True
                )
                
                with gr.Row():
                    btn_undo = gr.Button("⏪ Undo Move", variant="secondary")
                    btn_auto = gr.Button("🤖 AI Self-Play Move", variant="primary")
                
                with gr.Accordion("New Game Setup", open=False):
                    play_black_checkbox = gr.Checkbox(label="Play as Black (AI moves first)", value=False)
                    btn_reset = gr.Button("🔄 Restart Game", variant="stop")
            
            with gr.Tab("AI Parameters"):
                simulations_slider = gr.Slider(
                    minimum=5, 
                    maximum=100, 
                    value=25, 
                    step=5, 
                    label="MCTS Simulations", 
                    info="More sims = better play, but slower thinking"
                )
                temp_slider = gr.Slider(
                    minimum=0.0, 
                    maximum=1.0, 
                    value=0.1, 
                    step=0.05, 
                    label="Search Temperature", 
                    info="Higher temperature = more varied/creative moves"
                )
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 📈 Live Analysis")
                    analysis_box = gr.Markdown("Make a move to start analysis...")
            
                with gr.Column():
                    gr.Markdown("### 📝 Move Log")
                    log_box = gr.Markdown("*No moves made yet.*")

    # Hook Event Handlers
    move_dropdown.change(
        fn=make_player_move,
        inputs=[move_dropdown, simulations_slider, temp_slider],
        outputs=[board_html, status_markdown, move_dropdown, log_box, analysis_box]
    )
    
    btn_reset.click(
        fn=handle_reset,
        inputs=[play_black_checkbox, simulations_slider, temp_slider],
        outputs=[board_html, status_markdown, move_dropdown, log_box, analysis_box]
    )
    
    btn_undo.click(
        fn=handle_undo,
        inputs=[],
        outputs=[board_html, status_markdown, move_dropdown, log_box, analysis_box]
    )
    
    btn_auto.click(
        fn=handle_ai_vs_ai,
        inputs=[simulations_slider, temp_slider],
        outputs=[board_html, status_markdown, move_dropdown, log_box, analysis_box]
    )

demo.launch(theme=gr.themes.Soft(), css=custom_css)