import subprocess
import chess
import torch
import torch.nn as nn
import numpy as np
import hashlib
import matplotlib.pyplot as plt
from sklearn.metrics import mutual_info_score

# -------------------------------
# 1. Telepathic Memory Module (STLPL)
# -------------------------------
class TelepathicMemory(nn.Module):
    def __init__(self, d_model=64, nhead=4, num_layers=3,
                 collapse_threshold=0.5, temp=0.6, anchor_hash="pi_anchor:e_anchor"):
        super().__init__()
        self.d_model = d_model
        self.num_squares = 64
        self.collapse_threshold = collapse_threshold
        self.temp = temp

        # Embedding for piece types (13 classes: empty + 6 black + 6 white)
        self.piece_embed = nn.Embedding(13, d_model)
        # Positional encoding for squares (8x8)
        self.pos_embed = nn.Parameter(torch.randn(1, 64, d_model) * 0.02)

        # Telepathic self‑attention blocks (full bidirectional)
        self.attn_blocks = nn.ModuleList([
            nn.MultiheadAttention(d_model, nhead, batch_first=True)
            for _ in range(num_layers)
        ])

        self.gamma = nn.Parameter(torch.tensor(0.1))

        # Inter‑universal anchor Ξ
        self.xi_anchor = self._compute_xi_anchor(anchor_hash, d_model)

    def _compute_xi_anchor(self, anchor_hash, d_model):
        md5 = hashlib.md5(anchor_hash.encode()).hexdigest()
        vec = torch.zeros(d_model)
        for i in range(min(d_model, len(md5)//2)):
            byte_val = int(md5[2*i:2*i+2], 16)
            vec[i] = byte_val / 255.0
        return vec / vec.norm()

    def forward(self, board_tensor, steps=10, return_history=False, soft_collapse=False):
        batch = board_tensor.shape[0]
        tokens = board_tensor.view(batch, -1)           # (batch, 64)
        x = self.piece_embed(tokens) + self.pos_embed   # (batch, 64, d_model)

        history = []
        collapse_probs = []

        for _ in range(steps):
            for attn in self.attn_blocks:
                attn_out, _ = attn(x, x, x)
                grad = attn_out - x
                x = x + self.gamma * grad
                
                # Collapse calculation
                xi = self.xi_anchor.to(x.device).view(1, 1, -1)
                divergence = torch.norm(x - xi, dim=-1)
                p = torch.sigmoid((divergence - self.collapse_threshold) / self.temp)
                
                if soft_collapse:
                    # Differentiable soft-collapse approximation
                    x = p.unsqueeze(-1) * xi + (1.0 - p.unsqueeze(-1)) * x
                else:
                    # Stochastic collapse
                    mask = torch.bernoulli(p).unsqueeze(-1)
                    x = mask * xi + (1.0 - mask) * x
                
                collapse_probs.append(p.mean())
            history.append(x.detach().clone() if not soft_collapse else x)

        if return_history:
            return x, history, [p.item() if torch.is_tensor(p) else p for p in collapse_probs]
        return x

    def compute_phi(self, board_tensor, steps=20):
        _, history, _ = self.forward(board_tensor, steps=steps, return_history=True, soft_collapse=False)
        if len(history) < 2:
            return 0.0
        H = torch.stack(history, dim=0).squeeze(1)   # (steps, 64, d_model)
        steps, N, D = H.shape
        half = N // 2
        X = H[:, :half, :].reshape(steps, -1).detach().cpu().numpy()
        Y = H[:, half:, :].reshape(steps, -1).detach().cpu().numpy()
        Xd = (X > 0).astype(int)
        Yd = (Y > 0).astype(int)
        I_whole = mutual_info_score(
            [f"{x}{y}" for x, y in zip(Xd[:-1].flatten(), Yd[:-1].flatten())],
            [f"{x}{y}" for x, y in zip(Xd[1:].flatten(), Yd[1:].flatten())]
        )
        I_X = mutual_info_score(Xd[:-1].flatten(), Xd[1:].flatten())
        I_Y = mutual_info_score(Yd[:-1].flatten(), Yd[1:].flatten())
        phi = max(0.0, I_whole - (I_X + I_Y))
        return phi

# -------------------------------
# 2. Convert chess.Board to piece index tensor
# -------------------------------
PIECE_TO_IDX = {
    None: 0,                     # empty
    chess.PAWN: 1, chess.KNIGHT: 2, chess.BISHOP: 3,
    chess.ROOK: 4, chess.QUEEN: 5, chess.KING: 6,
}

def board_to_tensor(board: chess.Board) -> torch.Tensor:
    tensor = torch.zeros(8, 8, dtype=torch.long)
    for square in chess.SQUARES:
        piece = board.piece_at(square)
        if piece is None:
            idx = 0
        else:
            base = PIECE_TO_IDX[piece.piece_type]
            if piece.color == chess.WHITE:
                idx = base + 6   # white: 7-12
            else:
                idx = base       # black: 1-6
        row, col = divmod(square, 8)
        tensor[row, col] = idx
    return tensor.unsqueeze(0)   # add batch dimension

# -------------------------------
# 3. Interfacing GNU Chess via UCI
# -------------------------------
class GNUChessController:
    def __init__(self, engine_path="gnuchess"):
        self.board = chess.Board()
        self.use_mock = False
        try:
            self.engine = subprocess.Popen(
                [engine_path, "--uci"],
                universal_newlines=True,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            # Initialise UCI protocol
            self._write("uci")
            self._read_until("uciok")
            self._write("isready")
            self._read_until("readyok")
        except FileNotFoundError:
            print(f"Warning: '{engine_path}' engine executable not found. Running in mock mode.")
            self.use_mock = True
            self.engine = None

    def _write(self, cmd):
        if self.use_mock:
            return
        self.engine.stdin.write(cmd + "\n")
        self.engine.stdin.flush()

    def _read_line(self):
        if self.use_mock:
            return ""
        return self.engine.stdout.readline().strip()

    def _read_until(self, keyword):
        if self.use_mock:
            return
        while True:
            line = self._read_line()
            if keyword in line:
                break

    def set_position(self, fen=None):
        if fen is None:
            self._write("position startpos")
            self.board.set_fen(chess.STARTING_FEN)
        else:
            self._write(f"position fen {fen}")
            self.board.set_fen(fen)

    def make_move(self, uci_move):
        self._write(f"position fen {self.board.fen()} moves {uci_move}")
        self.board.push(chess.Move.from_uci(uci_move))

    def get_evaluation(self, thinking_time=0.2):
        if self.use_mock:
            return 10
        self._write(f"go movetime {int(thinking_time*1000)}")
        score = 0
        while True:
            line = self._read_line()
            if line.startswith("info"):
                parts = line.split()
                if "score" in parts:
                    idx = parts.index("score")
                    if idx + 2 < len(parts):
                        score_type = parts[idx+1]
                        score_val = parts[idx+2]
                        try:
                            score = int(score_val)
                            if score_type == "mate":
                                score = 1000 if score > 0 else -1000
                        except ValueError:
                            pass
            elif line.startswith("bestmove"):
                break
        return score

    def get_evaluation_and_best_move(self, thinking_time=0.5):
        if self.use_mock:
            import random
            if not self.board.legal_moves:
                return 10, None
            return 10, random.choice(list(self.board.legal_moves)).uci()

        self._write(f"go movetime {int(thinking_time*1000)}")
        best_move = None
        score = 0
        while True:
            line = self._read_line()
            if line.startswith("info"):
                parts = line.split()
                if "score" in parts:
                    idx = parts.index("score")
                    if idx + 2 < len(parts):
                        score_type = parts[idx+1]
                        score_val = parts[idx+2]
                        try:
                            score = int(score_val)
                            if score_type == "mate":
                                score = 1000 if score > 0 else -1000
                        except ValueError:
                            pass
            elif line.startswith("bestmove"):
                parts = line.split()
                best_move = parts[1] if len(parts) > 1 else None
                break
        return score, best_move

    def close(self):
        if self.use_mock:
            return
        self._write("quit")
        self.engine.terminate()

# -------------------------------
# 4. Rendering CLI Board
# -------------------------------
def print_board(board: chess.Board, player_color: chess.Color):
    flipped = (player_color == chess.BLACK)
    unicode_str = board.unicode(flipped=flipped, empty_square='·')
    lines = unicode_str.split('\n')
    
    ranks = range(8, 0, -1) if not flipped else range(1, 9)
    files_header = "    a b c d e f g h" if not flipped else "    h g f e d c b a"
    
    print(f"\n{files_header}")
    print("  +-----------------+")
    for rank, line in zip(ranks, lines):
        print(f"{rank} | {line} | {rank}")
    print("  +-----------------+")
    print(f"{files_header}\n")

# -------------------------------
# 5. Game Loop & Optimization
# -------------------------------
def main():
    print("Starting GNU Chess engine...")
    engine = GNUChessController()
    board = chess.Board()

    # Choose color
    color_input = input("Choose your color (W/B) [default: W]: ").strip().lower()
    player_color = chess.BLACK if color_input == 'b' else chess.WHITE
    
    # 2. Setup Telepathic Memory and Optimizer
    telepathic = TelepathicMemory(d_model=64, nhead=4, num_layers=3,
                                  collapse_threshold=0.5, temp=0.6)
    
    telepathic.train()
    optimizer = torch.optim.Adam(telepathic.parameters(), lr=0.005)

    # Setup Matplotlib Interactive Plot
    plt.ion()
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 6))
    fig.canvas.manager.set_window_title("Telepathic Chess Memory Alignment Dashboard")
    
    coherences = []
    target_coherences = []
    losses = []
    phis = []

    def optimize_memory(board_state, eval_score):
        # Differentiable training step
        optimizer.zero_grad()
        
        board_tensor = board_to_tensor(board_state)
        final_state, _, _ = telepathic.forward(board_tensor, steps=10, return_history=True, soft_collapse=True)
        xi = telepathic.xi_anchor.view(1, 1, -1)
        divergence_final = torch.norm(final_state - xi, dim=-1).mean()
        coherence = 1.0 / (1.0 + divergence_final)
        
        # Target coherence based on chess evaluation score
        # Map centipawns to [0.1, 0.9] range
        target_coherence = 0.1 + 0.8 * torch.sigmoid(torch.tensor(eval_score / 150.0))
        
        loss = (coherence - target_coherence) ** 2
        loss.backward()
        optimizer.step()
        
        phi = telepathic.compute_phi(board_tensor, steps=10)
        
        # Save metrics
        coherences.append(coherence.item())
        target_coherences.append(target_coherence.item())
        losses.append(loss.item())
        phis.append(phi)
        
        # Update plot
        ax1.clear()
        ax1.plot(coherences, label="Memory Coherence", color="#00ffd0", linewidth=2)
        ax1.plot(target_coherences, label="Target (Chess Metric)", color="#ff5555", linestyle="--", linewidth=1.5)
        ax1.set_title("Stationary Memory Coherence vs. Chess Metric")
        ax1.set_ylabel("Coherence")
        ax1.legend(loc="upper left")
        ax1.grid(True, linestyle=":", alpha=0.6)
        
        ax2.clear()
        ax2.plot(losses, label="Alignment Loss", color="#ffaa00", linewidth=2)
        ax2.plot(phis, label="Integrated Information Φ", color="#00ff00", linewidth=1.5)
        ax2.set_title("Optimization Metrics")
        ax2.set_xlabel("Game Moves")
        ax2.set_ylabel("Values")
        ax2.legend(loc="upper left")
        ax2.grid(True, linestyle=":", alpha=0.6)
        
        plt.tight_layout()
        plt.draw()
        plt.pause(0.1)
        
        print(f"📊 Coherence: {coherence.item():.3f} (Target: {target_coherence.item():.3f})")
        print(f"📈 Loss: {loss.item():.4f} | Integrated Info Φ: {phi:.4f}")

    print("\nGame Started!")
    print("Commands: Enter standard algebraic notation (e.g. e4, Nf3) or UCI notation (e.g. e2e4).")
    print("Use 'undo' to step back, or 'quit' to exit.")

    try:
        while not board.is_game_over():
            print_board(board, player_color)
            
            # Active side
            if board.turn == player_color:
                # Player Turn
                move_str = input("Your move: ").strip()
                if move_str.lower() == 'quit':
                    break
                if move_str.lower() == 'undo':
                    if len(board.move_stack) >= 2:
                        board.pop()
                        board.pop()
                        engine.set_position(board.fen())
                        if len(coherences) > 0:
                            coherences.pop()
                            target_coherences.pop()
                            losses.pop()
                            phis.pop()
                        print("Undid the last move pair.")
                    else:
                        print("Cannot undo (not enough moves played).")
                    continue
                
                try:
                    # Try SAN, then UCI
                    try:
                        move = board.parse_san(move_str)
                    except ValueError:
                        move = board.parse_uci(move_str)
                        
                    if move in board.legal_moves:
                        board.push(move)
                        engine.set_position(board.fen())
                    else:
                        print("❌ Illegal move. Try again.")
                        continue
                except ValueError:
                    print("❌ Invalid move format. Use SAN (e.g., e4, Nf3) or UCI (e.g., e2e4).")
                    continue
                
                # After player move, compute evaluation & train memory
                print("Computing stationary memory alignment...")
                score = engine.get_evaluation(thinking_time=0.3)
                optimize_memory(board, score)
                
            else:
                # Engine Turn
                print("GNU Chess thinking...")
                score, best_move = engine.get_evaluation_and_best_move(thinking_time=0.6)
                if best_move:
                    move = chess.Move.from_uci(best_move)
                    board.push(move)
                    engine.set_position(board.fen())
                    print(f"🤖 GNU Chess played: {board.san(move)}")
                    
                    # After engine move, train memory
                    print("Computing stationary memory alignment...")
                    optimize_memory(board, score)
                else:
                    break

        if board.is_game_over():
            print_board(board, player_color)
            print("\nGame Over!")
            print(f"Result: {board.result()}")
            
    except KeyboardInterrupt:
        print("\nExiting game...")
    finally:
        engine.close()
        plt.ioff()
        plt.show()

if __name__ == "__main__":
    main()