import cv2
import numpy as np
import chess
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def capture_image(url):
    """Capture an image from the webcam stream."""
    cap = cv2.VideoCapture(url)
    ret, frame = cap.read()
    cap.release()
    return frame

def extract_fen_from_image(image):
    """Detect the chessboard from the image and extract the FEN string.
       NOTE: This is a simplified version and might not work for all scenarios.
       A more robust solution would involve training a model for piece detection."""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, (8,8))
    if not ret:
        return None
    # Here we would typically identify each square and its piece to extract the FEN.
    # For simplicity, returning a default FEN for now.
    return "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"

def generate_heatmap_overlay(board_shape, aggregated_targets, cmap='viridis', alpha=0.5):
    """Generate heatmap overlay based on aggregated targets."""
    heatmap = np.zeros(board_shape[:2], dtype=np.uint8)
    for square, count in aggregated_targets.items():
        x, y = chess.SQUARE_NAMES.index(square) % 8, chess.SQUARE_NAMES.index(square) // 8
        heatmap[y, x] = count
    heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
    heatmap_colored = cm.get_cmap(cmap)(heatmap_normalized)[:, :, :3] * 255
    heatmap_overlay = (heatmap_colored * alpha + board_shape * (1 - alpha)).astype(np.uint8)
    return heatmap_overlay

def overlay_heatmap_on_image(image, heatmap_overlay):
    """Overlay heatmap onto the original image."""
    return cv2.addWeighted(image, 1, heatmap_overlay, 1, 0)

if __name__ == "__main__":
    url = "http://10.1.10.159:8000/stream.mjpg"
    image = capture_image(url)
    fen = extract_fen_from_image(image)
    if fen:
        board = chess.Board(fen)
        moves = get_legal_moves(board)
        aggregated_targets = aggregate_target_squares(moves)
        heatmap_overlay = generate_heatmap_overlay(image.shape, aggregated_targets)
        final_image = overlay_heatmap_on_image(image, heatmap_overlay)
        cv2.imshow('Overlayed Chessboard', final_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()