"""
Transposition Table for Nexus-Nano
64MB compact cache for speed
"""

import chess
import numpy as np
from typing import Optional, Dict, Tuple
from enum import Enum


class NodeType(Enum):
    EXACT = 0
    LOWER_BOUND = 1
    UPPER_BOUND = 2


class TTEntry:
    __slots__ = ['zobrist_key', 'depth', 'score', 'node_type', 'best_move', 'age']
    
    def __init__(self, zobrist_key, depth, score, node_type, best_move, age):
        self.zobrist_key = zobrist_key
        self.depth = depth
        self.score = score
        self.node_type = node_type
        self.best_move = best_move
        self.age = age


class TranspositionTable:
    """Compact 64MB transposition table"""
    
    def __init__(self, size_mb: int = 64):
        """Initialize with 64MB size"""
        
        bytes_per_entry = 64
        self.max_entries = (size_mb * 1024 * 1024) // bytes_per_entry
        
        self.table: Dict[int, TTEntry] = {}
        
        self.hits = 0
        self.misses = 0
        self.collisions = 0
        self.current_age = 0
        
        self._init_zobrist_keys()
    
    def _init_zobrist_keys(self):
        """Initialize Zobrist keys"""
        np.random.seed(42)
        
        self.zobrist_pieces = np.random.randint(
            0, 2**63, size=(12, 64), dtype=np.int64
        )
        self.zobrist_turn = np.random.randint(0, 2**63, dtype=np.int64)
        self.zobrist_castling = np.random.randint(0, 2**63, size=4, dtype=np.int64)
        self.zobrist_ep = np.random.randint(0, 2**63, size=8, dtype=np.int64)
    
    def compute_zobrist_key(self, board: chess.Board) -> int:
        """Fast Zobrist hash computation"""
        
        key = 0
        
        piece_to_idx = {
            (chess.PAWN, chess.WHITE): 0, (chess.KNIGHT, chess.WHITE): 1,
            (chess.BISHOP, chess.WHITE): 2, (chess.ROOK, chess.WHITE): 3,
            (chess.QUEEN, chess.WHITE): 4, (chess.KING, chess.WHITE): 5,
            (chess.PAWN, chess.BLACK): 6, (chess.KNIGHT, chess.BLACK): 7,
            (chess.BISHOP, chess.BLACK): 8, (chess.ROOK, chess.BLACK): 9,
            (chess.QUEEN, chess.BLACK): 10, (chess.KING, chess.BLACK): 11,
        }
        
        for square, piece in board.piece_map().items():
            idx = piece_to_idx[(piece.piece_type, piece.color)]
            key ^= self.zobrist_pieces[idx, square]
        
        if board.turn == chess.BLACK:
            key ^= self.zobrist_turn
        
        if board.has_kingside_castling_rights(chess.WHITE):
            key ^= self.zobrist_castling[0]
        if board.has_queenside_castling_rights(chess.WHITE):
            key ^= self.zobrist_castling[1]
        if board.has_kingside_castling_rights(chess.BLACK):
            key ^= self.zobrist_castling[2]
        if board.has_queenside_castling_rights(chess.BLACK):
            key ^= self.zobrist_castling[3]
        
        if board.ep_square is not None:
            key ^= self.zobrist_ep[board.ep_square % 8]
        
        return key
    
    def probe(self, zobrist_key, depth, alpha, beta) -> Optional[Tuple]:
        """Fast TT probe"""
        
        entry = self.table.get(zobrist_key)
        
        if entry is None:
            self.misses += 1
            return None
        
        if entry.zobrist_key != zobrist_key:
            self.collisions += 1
            return None
        
        if entry.depth < depth:
            self.misses += 1
            return None
        
        self.hits += 1
        
        score = entry.score
        
        if entry.node_type == NodeType.EXACT:
            return (score, entry.best_move)
        elif entry.node_type == NodeType.LOWER_BOUND and score >= beta:
            return (score, entry.best_move)
        elif entry.node_type == NodeType.UPPER_BOUND and score <= alpha:
            return (score, entry.best_move)
        
        return (None, entry.best_move)
    
    def store(self, zobrist_key, depth, score, node_type, best_move):
        """Store entry"""
        
        existing = self.table.get(zobrist_key)
        
        if existing and depth < existing.depth and existing.age == self.current_age:
            return
        
        self.table[zobrist_key] = TTEntry(
            zobrist_key, depth, score, node_type, best_move, self.current_age
        )
        
        if len(self.table) > self.max_entries:
            self._cleanup()
    
    def _cleanup(self):
        """Quick cleanup (10%)"""
        remove_count = self.max_entries // 10
        old_keys = sorted(self.table.keys(), key=lambda k: self.table[k].age)[:remove_count]
        for key in old_keys:
            del self.table[key]
    
    def increment_age(self):
        self.current_age += 1
    
    def clear(self):
        self.table.clear()
        self.hits = self.misses = self.collisions = 0
    
    def get_stats(self) -> Dict:
        total = self.hits + self.misses
        hit_rate = (self.hits / total * 100) if total > 0 else 0
        
        return {
            'entries': len(self.table),
            'max_entries': self.max_entries,
            'usage_percent': len(self.table) / self.max_entries * 100,
            'hits': self.hits,
            'misses': self.misses,
            'hit_rate': hit_rate,
            'collisions': self.collisions
        }