import chess.svg
import chess
import re
import cv2
import random as rd
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import load_model
import chess, chess.svg
from pathlib import Path
import webbrowser
piece_symbols = 'prbnkqPRBNKQ-'
try:
    model = load_model('./chess_model.h5')
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    
def onehot_from_fen(fen):
    eye = np.eye(13)
    output = np.empty((0, 13))
    fen = re.sub('[-]', '', fen)

    for char in fen:
        if(char in '12345678'):
            output = np.append(output, np.tile(eye[12], (int(char), 1)), axis=0)
        else:
            idx = piece_symbols.index(char)
            output = np.append(output, eye[idx].reshape((1, 13)), axis=0)
    return output
import numpy as np
import re



def board_from_fen(fen):
    board = []
    fen = fen.split()[0]   # just the board part (ignore turn, castling, etc.)
    
    for row in fen.split('/'):
        row_out = []
        for char in row:
            if char.isdigit():
                # expand empty squares
                row_out.extend([0] * int(char))
            else:
                idx = piece_symbols.index(char) + 1  # +1 so 0 = empty
                row_out.append(idx)
        board.append(row_out)

    return np.array(board, dtype=np.int8)  # shape (8,8)

def fen_from_onehot(one_hot):
    output = ''
    for j in range(8):
        for i in range(8):
            if(one_hot[j][i] == 12):
                output += ' '
            else:
                output += piece_symbols[one_hot[j][i]]
        if(j != 7):
            output += '-'

    for i in range(8, 0, -1):
        output = output.replace(' ' * i, str(i))
    return output


def order_points(pts):
    # Order 4 points: top-left, top-right, bottom-right, bottom-left
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect

def preprocess_chessboard(image_path, target_size=(30, 30)):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5,5), 0)
    edges = cv2.Canny(blur, 50, 150)

    # Find contours
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    cropped_board = None
    warped = None
    squares = []

    if contours:
        cnt = max(contours, key=cv2.contourArea)
        epsilon = 0.02 * cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, epsilon, True)

        # If we found 4 corners, warp
        if len(approx) == 4:
            pts = approx.reshape(4,2)
            rect = order_points(pts)
            board_size = 256  # fixed square board
            dst = np.array([
                [0, 0],
                [board_size-1, 0],
                [board_size-1, board_size-1],
                [0, board_size-1]
            ], dtype="float32")

            M = cv2.getPerspectiveTransform(rect, dst)
            warped = cv2.warpPerspective(img, M, (board_size, board_size))
            cropped_board = warped
        else:
            # fallback: just resize whole image
            warped = cv2.resize(img, (256, 256))
            cropped_board = warped
    else:
        # no contours, fallback
        warped = cv2.resize(img, (256, 256))
        cropped_board = warped

    # Split into 64 squares
    board_size = warped.shape[0]
    sq_size = board_size // 8
    for row in range(8):
        for col in range(8):
            square = warped[row*sq_size:(row+1)*sq_size,
                            col*sq_size:(col+1)*sq_size]
            square = cv2.resize(square, target_size)
            square = square / 255.0
            squares.append(square)

    return np.array(squares), cropped_board

def display_with_predicted_fen(image):
    squares,_ = preprocess_chessboard(image)
    pred = model.predict(squares).argmax(axis=1).reshape(8, 8)
    fen = fen_from_onehot(pred)
    return fen

def get_board_from_image(img_path):
    image = cv2.imread(img_path)
    if image is None:
        print('sorry image is null')
        return
    ext = img_path.split('.')[-1]
    if not (ext=='jpg' or ext == 'jpeg' or ext == 'png'):
        print('sorry image needs to be in jpg/jpeg/png format')
        return
    predicted_fen = display_with_predicted_fen(img_path)
    # Create board from predicted FEN
    board = chess.Board(predicted_fen.replace('-','/'))
    return predicted_fen
def get_board_from_fen(fen:str):
    print(type(fen))
    fen=fen.split('-')
    if len(fen)!=8:
        return [[]]
    answers = [[] for _ in range(8)]
    map={'p':'bp','r':'bR','n':'bN','k':'bK','q':'bQ','b':'bB','P':'wp','R':'wR','N':'wN','K':'wK','Q':'wQ','B':'wB'}
    for i,f in enumerate(fen):
        for char in f:
            
            if char in '12345678':
                answers[i].extend(['--']*(int(char)))
            elif char in piece_symbols:
                answers[i].append(map[char])
    return answers
                
    
def get_board(image_path):
    fen_notation = get_board_from_image(image_path)
    if fen_notation:
        board = get_board_from_fen(fen_notation)
        
    else:
        return [[]],""
    return board,fen_notation.replace('-','/')
    
if __name__ == "__main__":
    get_board()