import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import cv2
import pyautogui
import mss
from PIL import Image
import hashlib
from collections import OrderedDict

# ------------------------------
# 1. Screen capture & board recognition
# ------------------------------
def capture_chessboard(region=None):
    """
    Capture the CoreChess board region.
    If region not given, prompt user to select window manually (simple approach).
    Returns: numpy array (BGR image)
    """
    with mss.mss() as sct:
        if region is None:
            # Default: try to find window by title (adjust for your OS)
            # For simplicity, we capture the whole screen and let the user crop later
            monitor = sct.monitors[1]  # primary monitor
            screenshot = sct.grab(monitor)
            img = np.array(screenshot)
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            # Let user select ROI (OpenCV popup)
            roi = cv2.selectROI("Select chessboard", img, False)
            cv2.destroyWindow("Select chessboard")
            x, y, w, h = roi
            board_img = img[y:y+h, x:x+w]
        else:
            board_img = np.array(sct.grab(region))
            board_img = cv2.cvtColor(board_img, cv2.COLOR_BGRA2BGR)
        return board_img

# ------------------------------
# 2. Piece detection using a lightweight CNN
# ------------------------------
class PieceClassifier(nn.Module):
    """Simple CNN to classify 13 categories: empty + 12 piece types (6 colors*2)"""
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.fc = nn.Sequential(
            nn.Linear(128 * 8 * 8, 256), nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(256, 13)   # 13 classes
        )

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# Mapping indices to piece strings (FEN style)
PIECES = ['empty', 'p', 'n', 'b', 'r', 'q', 'k',  # black pieces (lowercase)
          'P', 'N', 'B', 'R', 'Q', 'K']            # white pieces (uppercase)
# We'll train the classifier on synthetic data – for demo we provide a placeholder.
# In practice, you would train on a dataset of cropped piece images.
def load_pretrained_classifier(weights_path=None):
    model = PieceClassifier()
    if weights_path:
        model.load_state_dict(torch.load(weights_path, map_location='cpu'))
    else:
        # Fallback: random weights (not usable). Provide a simple rule-based detection as placeholder.
        print("Warning: No trained classifier weights. Using dummy detection (all empty).")
        model.eval()
    return model

def detect_pieces(board_img, classifier, square_size=64):
    """
    board_img: RGB image of chessboard
    square_size: approximate pixel size of one square (assumed square).
    Returns 8x8 list of piece indices (0=empty, 1..12 piece types)
    """
    h, w = board_img.shape[:2]
    rows = cols = 8
    square_h = h // rows
    square_w = w // cols
    board_state = np.zeros((rows, cols), dtype=int)

    for i in range(rows):
        for j in range(cols):
            # crop square
            y1, y2 = i * square_h, (i+1) * square_h
            x1, x2 = j * square_w, (j+1) * square_w
            square = board_img[y1:y2, x1:x2]
            # preprocess
            square_resized = cv2.resize(square, (64, 64))
            square_tensor = torch.from_numpy(square_resized).permute(2,0,1).float() / 255.0
            square_tensor = square_tensor.unsqueeze(0)
            with torch.no_grad():
                logits = classifier(square_tensor)
                pred = torch.argmax(logits, dim=1).item()
            board_state[i, j] = pred
    return board_state

# ------------------------------
# 3. Telepathic Memory Module (STLPL core)
# ------------------------------
class TelepathicMemory(nn.Module):
    """
    Implements a self-telepathic loop over the 64 squares.
    Each square is a token with initial embedding = one-hot piece type.
    Recurrent self-attention with stochastic collapse to anchor Ξ.
    """
    def __init__(self, d_model=128, nhead=8, num_layers=4, 
                 collapse_threshold=0.5, temp=0.8, 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 type (13 categories)
        self.piece_embed = nn.Embedding(13, d_model)
        # Positional encoding for square coordinates (i,j)
        self.pos_embed = nn.Parameter(torch.randn(1, 64, d_model) * 0.02)

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

        # Learnable update gate (like ADDP)
        self.gamma = nn.Parameter(torch.tensor(0.1))

        # Collapse layers: each block has its own (we share parameters)
        self.xi_anchor = self._compute_xi_anchor(anchor_hash, d_model)
        self.collapse = self._make_collapse_layer()

    def _compute_xi_anchor(self, anchor_hash, d_model):
        """Derive Ξ from MD5 of given string."""
        md5 = hashlib.md5(anchor_hash.encode()).hexdigest()
        # Convert first d_model bytes to normalized float vector
        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 _make_collapse_layer(self):
        """Stochastic collapse function."""
        def collapse(h):
            # h: (batch, seq, d_model)
            xi = self.xi_anchor.to(h.device).view(1,1,-1)
            divergence = torch.norm(h - xi, dim=-1)  # (batch, seq)
            p_collapse = torch.sigmoid((divergence - self.collapse_threshold) / self.temp)
            mask = torch.bernoulli(p_collapse).unsqueeze(-1)  # (batch, seq, 1)
            h_collapsed = mask * xi + (1 - mask) * h
            return h_collapsed, p_collapse
        return collapse

    def forward(self, board_state, steps=10, return_history=False):
        """
        board_state: (batch, 8, 8) integer piece indices (0..12)
        returns: final hidden states (batch, 64, d_model)
        """
        batch = board_state.shape[0]
        # Flatten to tokens
        tokens = board_state.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:
                # Self-attention (telepathic reading)
                attn_out, _ = attn(x, x, x)
                # ADDP: update with gradient (difference between attended and original)
                grad = attn_out - x
                x = x + self.gamma * grad
                # Collapse
                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_state, steps=20):
        """
        Approximate integrated information Φ by splitting the 64 squares
        into two halves (first 32 vs last 32) and computing time‑lagged
        mutual information across the telepathic loop.
        """
        x, history, _ = self.forward(board_state, steps=steps, return_history=True)
        # history: list of (batch, 64, d_model)
        if len(history) < 2:
            return 0.0
        # Take first batch (batch=1 for simplicity)
        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)
        from sklearn.metrics import mutual_info_score
        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

# ------------------------------
# 4. Main pipeline: capture → detect → telepathic stationary memory
# ------------------------------
def main():
    print("Step 1: Capture CoreChess window...")
    board_img = capture_chessboard()
    cv2.imwrite("captured_board.png", board_img)
    print("Saved board image as 'captured_board.png'")

    # Load piece classifier (you need to train or provide weights)
    # For demonstration, we create a dummy classifier that always returns "empty"
    # Replace with your trained model.
    classifier = PieceClassifier()
    # Dummy forward that always predicts class 0 (empty) – for demo only.
    def dummy_forward(x): return torch.zeros((x.size(0), 13))
    classifier.forward = dummy_forward

    print("Step 2: Detecting pieces (dummy detection – all empty)...")
    board_state = detect_pieces(board_img, classifier)
    print("Detected board (indices):\n", board_state)

    # Convert to torch tensor
    board_tensor = torch.tensor(board_state, dtype=torch.long).unsqueeze(0)  # batch=1

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

    print("Step 3: 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}")

    # Interpret result: final collapsed representation can be used to evaluate the position.
    # For instance, compute a scalar "coherence" score:
    coherence = torch.sigmoid(final_state.mean()).item()
    print(f"Stationary coherence = {coherence:.3f} (higher = more self‑telepathically stable)")

    # Optional: visualize collapse probability over steps
    import matplotlib.pyplot as plt
    plt.plot(collapse_probs)
    plt.xlabel("Telepathic step")
    plt.ylabel("Average collapse probability")
    plt.title("Collapse dynamics")
    plt.show()

if __name__ == "__main__":
    main()