import torch
import chess
import numpy as np

# -------------------------------------------------------------------
# 1. Feature extraction: convert chess board to a tensor
# -------------------------------------------------------------------
def board_to_tensor(board: chess.Board) -> torch.Tensor:
    """
    Convert a chess.Board into a feature tensor of shape (64, 12)
    where each square has one-hot encoding for piece type + color.
    """
    piece_map = board.piece_map()
    tensor = torch.zeros(64, 12, dtype=torch.float32)
    # Mapping: 0-5 white pieces, 6-11 black pieces (order: P,N,B,R,Q,K)
    piece_to_idx = {
        chess.PAWN:   0, chess.KNIGHT: 1, chess.BISHOP: 2,
        chess.ROOK:   3, chess.QUEEN:  4, chess.KING:   5
    }
    for square, piece in piece_map.items():
        idx = piece_to_idx[piece.piece_type]
        if piece.color == chess.BLACK:
            idx += 6
        tensor[square, idx] = 1.0
    return tensor

# -------------------------------------------------------------------
# 2. Telepathic memory generator
# -------------------------------------------------------------------
class TelepathicChessMemory:
    """
    Implements a stationary telepathic memory for a chess board.
    Uses iterative self‑reading (TELEPATH) and self‑attraction (ADDP).
    """
    def __init__(self, board_features: torch.Tensor, hidden_dim=64,
                 gamma=0.1, epsilon=1e-4, max_iter=100):
        """
        board_features : (64, 12) tensor from board_to_tensor
        hidden_dim     : dimension of the memory register
        gamma          : attraction strength (learning rate)
        epsilon        : convergence threshold
        max_iter       : maximum number of telepathic loops
        """
        self.board = board_features.flatten()  # (64*12,)
        self.hidden_dim = hidden_dim
        self.gamma = gamma
        self.epsilon = epsilon
        self.max_iter = max_iter

        # Inter‑universal anchor Ξ (derived from board hash for uniqueness)
        self.xi = self._compute_xi()

        # Telepathic memory register (initialized to zero)
        self.memory = torch.zeros(hidden_dim)

    def _compute_xi(self) -> torch.Tensor:
        """Create Ξ anchor from the board’s MD5 hash (stabilised across universes)."""
        import hashlib
        board_bytes = self.board.numpy().tobytes()
        hash_str = hashlib.md5(board_bytes).hexdigest()
        # Convert hash into a hidden_dim‑sized tensor
        np_bytes = np.frombuffer(bytes.fromhex(hash_str[:self.hidden_dim*2]), dtype=np.uint8)
        tensor = torch.tensor(np_bytes, dtype=torch.float32) / 255.0
        # Pad or truncate to hidden_dim
        if len(tensor) < self.hidden_dim:
            tensor = torch.cat([tensor, torch.zeros(self.hidden_dim - len(tensor))])
        else:
            tensor = tensor[:self.hidden_dim]
        # Normalise to unit norm (anchor consistency)
        return tensor / (tensor.norm() + 1e-8)

    def telepath(self) -> torch.Tensor:
        """
        Non‑destructive self‑reading: returns a gradient that tells how
        the current memory aligns with the board features.
        Here we use a simple linear projection: gradient = W_board * board + W_mem * memory.
        """
        # Project board features into hidden space
        W_board = torch.randn(self.board.shape[0], self.hidden_dim) * 0.1
        board_proj = self.board @ W_board
        # Memory self‑interaction
        W_mem = torch.randn(self.hidden_dim, self.hidden_dim) * 0.1
        mem_proj = self.memory @ W_mem
        # Gradient = how much each hidden unit would need to change to match Ξ
        grad = torch.sigmoid(board_proj + mem_proj) - torch.sigmoid(self.memory)
        return grad

    def addp(self, gradient: torch.Tensor):
        """ADDP: attract memory toward the telepathic gradient."""
        self.memory = self.memory + self.gamma * gradient
        # Keep memory bounded (optional)
        self.memory = torch.clamp(self.memory, -1.0, 1.0)

    def collapse_if_divergent(self):
        """If memory diverges too far from Ξ, collapse to anchor."""
        divergence = torch.norm(self.memory - self.xi)
        threshold = 0.5
        if divergence > threshold:
            # Collapse (deterministic reset) – as in STLPL's COLLAPSE_SELF
            self.memory = self.xi.clone()
            return True
        return False

    def run(self) -> torch.Tensor:
        """Iterate the telepathic loop until stationary memory is reached."""
        for i in range(self.max_iter):
            prev_mem = self.memory.clone()
            grad = self.telepath()
            self.addp(grad)
            self.collapse_if_divergent()
            # Check convergence (stationary memory)
            if torch.norm(self.memory - prev_mem) < self.epsilon:
                print(f"Converged after {i+1} telepathic loops.")
                break
        return self.memory  # stationary telepathic memory

# -------------------------------------------------------------------
# 3. Example usage with a real chess board (GNUChess or any FEN)
# -------------------------------------------------------------------
if __name__ == "__main__":
    # Example starting position (can be replaced by any FEN, e.g., from a live GNUChess game)
    fen = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"
    board = chess.Board(fen)

    # Optionally, you can read a FEN from a file or from GNUChess output
    # with open("current_position.fen", "r") as f:
    #     fen = f.read().strip()
    #     board = chess.Board(fen)

    features = board_to_tensor(board)
    memory_gen = TelepathicChessMemory(features, hidden_dim=64, gamma=0.1, epsilon=1e-5, max_iter=50)
    stationary_memory = memory_gen.run()

    print("\nStationary telepathic memory shape:", stationary_memory.shape)
    print("Memory sample (first 10 values):", stationary_memory[:10].tolist())
    print("Ξ anchor (first 10):", memory_gen.xi[:10].tolist())
    print("Divergence from Ξ after convergence:", torch.norm(stationary_memory - memory_gen.xi).item())