"""
Blackwell Knight - ChessNet Modeling File
Custom chess neural network for the TobiasLogic Chess Arena.

Architecture: ResNet backbone with policy + value heads.
Board input: 19 planes of 8x8.
Move output: Flat index over 20,480 possible moves.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import chess

# =====================================================================
# BOARD ENCODING
# =====================================================================

def encode_board(board):
    """
    Encode a chess.Board into a (19, 8, 8) float32 tensor.
    
    Planes:
      0-5:   White pieces (Pawn, Knight, Bishop, Rook, Queen, King)
      6-11:  Black pieces (Pawn, Knight, Bishop, Rook, Queen, King)
      12:    Side to move (1.0 = white, 0.0 = black)
      13-16: Castling rights (WK, WQ, BK, BQ)
      17:    En passant target square
      18:    Halfmove clock (normalized by /100)
    """
    planes = np.zeros((19, 64), dtype=np.float32)

    for square, piece in board.piece_map().items():
        idx = piece.piece_type - 1 + (0 if piece.color == chess.WHITE else 6)
        planes[idx, square] = 1.0

    planes[12].fill(1.0 if board.turn == chess.WHITE else 0.0)
    planes[13].fill(1.0 if board.has_kingside_castling_rights(chess.WHITE) else 0.0)
    planes[14].fill(1.0 if board.has_queenside_castling_rights(chess.WHITE) else 0.0)
    planes[15].fill(1.0 if board.has_kingside_castling_rights(chess.BLACK) else 0.0)
    planes[16].fill(1.0 if board.has_queenside_castling_rights(chess.BLACK) else 0.0)

    if board.ep_square is not None:
        planes[17, board.ep_square] = 1.0

    planes[18].fill(min(board.halfmove_clock / 100.0, 1.0))

    return planes.reshape(19, 8, 8)


# =====================================================================
# MOVE ENCODING / DECODING
# =====================================================================

POLICY_SIZE = 20480  # 64 * 64 * 5

def get_move_index(move):
    """Convert a chess.Move to a flat index (0-20479)."""
    promo_idx = 0
    if move.promotion:
        promo_idx = {
            chess.KNIGHT: 1,
            chess.BISHOP: 2,
            chess.ROOK: 3,
            chess.QUEEN: 4,
        }.get(move.promotion, 0)
    return (move.from_square * 64 + move.to_square) * 5 + promo_idx


def index_to_move(idx, board):
    """
    Convert a flat index back to a chess.Move.
    If the decoded move is illegal, returns None.
    """
    from_sq = (idx // 5) // 64
    to_sq = (idx // 5) % 64
    promo_idx = idx % 5

    promo = None
    if promo_idx == 1:
        promo = chess.KNIGHT
    elif promo_idx == 2:
        promo = chess.BISHOP
    elif promo_idx == 3:
        promo = chess.ROOK
    elif promo_idx == 4:
        promo = chess.QUEEN

    move = chess.Move(from_sq, to_sq, promotion=promo)

    # Safety: check if the move is actually legal
    if move in board.legal_moves:
        return move
    return None


# =====================================================================
# MODEL ARCHITECTURE
# =====================================================================

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        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):
        identity = x
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += identity
        return F.relu(out)


class ChessNet(nn.Module):
    """
    Blackwell Knight ChessNet.
    
    ~18.5M parameters.
    Input:  (batch, 19, 8, 8) board planes.
    Output: policy_logits (batch, 20480), value (batch, 1) in [-1, 1].
    """

    def __init__(self, config=None):
        super().__init__()

        if config is None:
            config = {}

        in_planes = config.get("input_planes", 19)
        num_blocks = config.get("num_blocks", 10)
        num_channels = config.get("num_channels", 128)
        policy_channels = config.get("policy_channels", 32)
        value_channels = config.get("value_channels", 8)
        policy_size = config.get("policy_size", POLICY_SIZE)

        self.stem = nn.Sequential(
            nn.Conv2d(in_planes, num_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(num_channels),
            nn.ReLU(),
        )

        self.blocks = nn.Sequential(
            *[ResidualBlock(num_channels) for _ in range(num_blocks)]
        )

        self.policy_head = nn.Sequential(
            nn.Conv2d(num_channels, policy_channels, 1, bias=False),
            nn.BatchNorm2d(policy_channels),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(policy_channels * 8 * 8, policy_size),
        )

        self.value_head = nn.Sequential(
            nn.Conv2d(num_channels, value_channels, 1, bias=False),
            nn.BatchNorm2d(value_channels),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(value_channels * 8 * 8, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Tanh(),
        )

    def forward(self, x):
        x = self.stem(x)
        x = self.blocks(x)
        policy_logits = self.policy_head(x)
        value = self.value_head(x)
        return policy_logits, value


# =====================================================================
# INFERENCE / PLAY CHESS
# =====================================================================

class BlackwellKnightPlayer:
    """
    Arena-compatible player wrapper.
    
    Usage:
        player = BlackwellKnightPlayer("akshit/blackwell-knight-v1")
        move = player.choose_move(board)
    """

    def __init__(self, model_repo, device=None, temperature=0.0):
        from huggingface_hub import hf_hub_download
        from safetensors.torch import load_file
        import json

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.temperature = temperature

        # Load config
        config_path = hf_hub_download(model_repo, "config.json")
        with open(config_path, "r") as f:
            self.config = json.load(f)

        # Build model
        arch = self.config.get("model_architecture", {})
        self.model = ChessNet({
            "input_planes": self.config.get("board_encoding", {}).get("input_planes", 19),
            "num_blocks": arch.get("num_blocks", 10),
            "num_channels": arch.get("num_channels", 128),
            "policy_channels": arch.get("policy_channels", 32),
            "value_channels": arch.get("value_channels", 8),
            "policy_size": self.config.get("move_encoding", {}).get("policy_size", POLICY_SIZE),
        })

        # Load weights
        weights_path = hf_hub_download(model_repo, "model.safetensors")
        state_dict = load_file(weights_path)
        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

    def choose_move(self, board):
        """
        Given a chess.Board, return the best legal move.
        """
        x = encode_board(board)
        x_tensor = torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(self.device)

        with torch.no_grad():
            policy_logits, value = self.model(x_tensor)

        # Get legal moves
        legal_moves = list(board.legal_moves)
        if not legal_moves:
            return None

        legal_indices = [get_move_index(m) for m in legal_moves]

        # Mask illegal moves
        mask = torch.full((POLICY_SIZE,), float("-inf"), device=self.device)
        mask[legal_indices] = 0.0

        logits = policy_logits[0] + mask

        if self.temperature > 0:
            probs = torch.softmax(logits / self.temperature, dim=-1)
            idx = torch.multinomial(probs, 1).item()
        else:
            idx = torch.argmax(logits).item()

        move = index_to_move(idx, board)

        # Fallback: if something went wrong, pick a random legal move
        if move is None or move not in legal_moves:
            import random
            move = random.choice(legal_moves)

        return move

    def evaluate(self, board):
        """
        Return the value estimate for the current position.
        Positive = good for side to move.
        """
        x = encode_board(board)
        x_tensor = torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(self.device)

        with torch.no_grad():
            _, value = self.model(x_tensor)

        return value.item()