import torch
import chess
import numpy as np
import hashlib

# ------------------------------------------------------------
# Feature extraction: chess board → tensor
# ------------------------------------------------------------
def board_to_tensor(board: chess.Board) -> torch.Tensor:
    """Convert board to a flat feature vector (64 squares × 12 piece types)."""
    piece_map = board.piece_map()
    tensor = torch.zeros(64, 12, dtype=torch.float32)
    piece_to_idx = {
        chess.PAWN: 0, chess.KNIGHT: 1, chess.BISHOP: 2,
        chess.ROOK: 3, chess.QUEEN: 4, chess.KING: 5
    }
    for sq, piece in piece_map.items():
        idx = piece_to_idx[piece.piece_type]
        if piece.color == chess.BLACK:
            idx += 6
        tensor[sq, idx] = 1.0
    return tensor.flatten()  # (768,)

# ------------------------------------------------------------
# Telepathic memory generator (interactive version)
# ------------------------------------------------------------
class TelepathicChessMemory:
    def __init__(self, board_features: torch.Tensor, hidden_dim=32,
                 gamma=0.2, epsilon=1e-4, max_iter=100):
        self.board = board_features
        self.hidden_dim = hidden_dim
        self.gamma = gamma
        self.epsilon = epsilon
        self.max_iter = max_iter

        # Random projection matrices (fixed for reproducibility)
        self.W_board = torch.randn(len(board_features), hidden_dim) * 0.1
        self.W_mem   = torch.randn(hidden_dim, hidden_dim) * 0.1

        # Inter‑universal anchor Ξ derived from the initial board
        self.xi = self._compute_xi()
        self.memory = torch.zeros(hidden_dim)

    def _compute_xi(self) -> torch.Tensor:
        """Stable anchor from board hash."""
        board_bytes = self.board.numpy().tobytes()
        h = hashlib.md5(board_bytes).hexdigest()
        # Convert first half of hash to hidden_dim floats
        n = self.hidden_dim
        bytes_needed = n * 4  # 4 bytes per float32
        hex_bytes = h[:bytes_needed*2]  # 2 hex chars per byte
        if len(hex_bytes) < bytes_needed*2:
            hex_bytes = hex_bytes.ljust(bytes_needed*2, '0')
        arr = np.frombuffer(bytes.fromhex(hex_bytes), dtype=np.float32)
        if len(arr) < n:
            arr = np.pad(arr, (0, n - len(arr)))
        else:
            arr = arr[:n]
        xi = torch.tensor(arr)
        return xi / (xi.norm() + 1e-8)

    def telepath(self) -> torch.Tensor:
        """Non‑destructive self‑reading: gradient that aligns memory with board."""
        board_proj = self.board @ self.W_board
        mem_proj   = self.memory @ self.W_mem
        # How much each memory unit would change to better match the board + itself
        target = torch.sigmoid(board_proj + mem_proj)
        grad = target - torch.sigmoid(self.memory)
        return grad

    def addp(self, grad: torch.Tensor):
        """Attract memory toward the telepathic gradient."""
        self.memory += self.gamma * grad
        self.memory = torch.clamp(self.memory, -1.0, 1.0)

    def collapse_if_divergent(self) -> bool:
        """If too far from Ξ, reset to anchor."""
        if torch.norm(self.memory - self.xi) > 0.8:
            self.memory = self.xi.clone()
            return True
        return False

    def run(self) -> torch.Tensor:
        """Iterate until stationary memory (fixed point)."""
        for _ in range(self.max_iter):
            prev = self.memory.clone()
            grad = self.telepath()
            self.addp(grad)
            self.collapse_if_divergent()
            if torch.norm(self.memory - prev) < self.epsilon:
                break
        return self.memory

# ------------------------------------------------------------
# Interactive chess loop
# ------------------------------------------------------------
def main():
    board = chess.Board()
    print("♜  Telepathic Chess Memory Interactive  ♞")
    print("Enter moves in SAN (e.g., e4, Nf3, O-O). Type 'quit' to exit.\n")

    # Compute initial stationary memory
    features = board_to_tensor(board)
    tmem = TelepathicChessMemory(features, hidden_dim=32, gamma=0.2)
    prev_memory = tmem.run()
    print("Initial memory generated. Divergence from Ξ:",
          torch.norm(tmem.memory - tmem.xi).item())

    while True:
        print("\n" + str(board))
        move_san = input("Your move: ").strip()
        if move_san.lower() == 'quit':
            break

        try:
            move = board.parse_san(move_san)
        except ValueError:
            print("Invalid move. Try again.")
            continue

        board.push(move)

        # Regenerate telepathic memory for new board
        features = board_to_tensor(board)
        tmem = TelepathicChessMemory(features, hidden_dim=32, gamma=0.2)
        new_memory = tmem.run()

        # --- Memory metric ---
        memory_change = torch.norm(new_memory - prev_memory).item()
        divergence_xi = torch.norm(new_memory - tmem.xi).item()

        print(f"\n📊 Memory metric after {move_san}:")
        print(f"   • Change from previous memory: {memory_change:.4f}")
        print(f"   • Distance to Ξ anchor:        {divergence_xi:.4f}")
        print(f"   • First 5 memory values:       {new_memory[:5].tolist()}")

        # Check for checkmate
        if board.is_checkmate():
            print("\n🏆 Checkmate! Game over.")
            break
        elif board.is_stalemate():
            print("\n♾️ Stalemate.")
            break

        prev_memory = new_memory

    print("Goodbye.")

if __name__ == "__main__":
    main()