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

# -------------------------------
# 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 _collapse(self, h):
        xi = self.xi_anchor.to(h.device).view(1, 1, -1)
        divergence = torch.norm(h - xi, dim=-1)
        p_collapse = torch.sigmoid((divergence - self.collapse_threshold) / self.temp)
        mask = torch.bernoulli(p_collapse).unsqueeze(-1)
        h_collapsed = mask * xi + (1 - mask) * h
        return h_collapsed, p_collapse

    def forward(self, board_tensor, steps=10, return_history=False):
        """
        board_tensor: (batch, 8, 8) piece indices (0 = empty,
                       1-6: black pieces P,N,B,R,Q,K; 7-12: white pieces in same order)
        """
        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
                x, p = self._collapse(x)
                collapse_probs.append(p.mean().item())
            history.append(x.detach().clone())

        if return_history:
            return x, history, collapse_probs
        return x

    def compute_phi(self, board_tensor, steps=20):
        _, history, _ = self.forward(board_tensor, steps=steps, return_history=True)
        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().numpy()
        Y = H[:, half:, :].reshape(steps, -1).detach().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,
}
# Black pieces get indices 1-6, white pieces get 7-12
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 (using random legal moves).")
            print("To use the real engine, please install it (e.g., 'sudo apt install gnuchess').")
            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_best_move(self, thinking_time=1.0):
        if self.use_mock:
            import random
            if not self.board.legal_moves:
                return None
            return random.choice(list(self.board.legal_moves)).uci()

        self._write(f"go movetime {int(thinking_time*1000)}")
        best_move = None
        while True:
            line = self._read_line()
            if line.startswith("bestmove"):
                parts = line.split()
                best_move = parts[1] if len(parts) > 1 else None
                break
        return best_move

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

# -------------------------------
# 4. Main: telepathic memory generation from GNU Chess
# -------------------------------
def main():
    print("Starting GNU Chess engine...")
    engine = GNUChessController()
    # Example: set up a position (starting position)
    engine.set_position()   # startpos
    print("Current FEN:", engine.board.fen())

    # Convert board to piece tensor
    board_tensor = board_to_tensor(engine.board)
    print("Board tensor shape:", board_tensor.shape)

    # Create telepathic memory module
    telepathic = TelepathicMemory(d_model=64, nhead=4, num_layers=3,
                                  collapse_threshold=0.5, temp=0.6)

    print("Running telepathic loop (stationary memory generation)...")
    final_state, history, collapse_probs = telepathic.forward(board_tensor,
                                                              steps=15,
                                                              return_history=True)

    # Compute integrated information Φ
    phi = telepathic.compute_phi(board_tensor, steps=15)
    print(f"Integrated information Φ = {phi:.4f}")

    # Coherence score: how close the final hidden state is to the anchor Ξ
    xi = telepathic.xi_anchor.view(1,1,-1)
    divergence_final = torch.norm(final_state - xi, dim=-1).mean().item()
    coherence = 1.0 / (1.0 + divergence_final)   # high when close to Ξ
    print(f"Stationary coherence = {coherence:.3f}")

    # Optional: Get GNU Chess's best move for comparison
    best_move_uci = engine.get_best_move(thinking_time=1.0)
    print(f"GNU Chess best move: {best_move_uci}")

    # Visualise collapse dynamics
    plt.plot(collapse_probs)
    plt.xlabel("Telepathic step")
    plt.ylabel("Average collapse probability")
    plt.title("Collapse dynamics during telepathic memory formation")
    plt.show()

    engine.close()

if __name__ == "__main__":
    main()