import cv2
import numpy as np
from detector import Detector
from board_manager import BoardManager

# -------------------- Pièces --------------------
def extract_pieces(pieces_pred):
    """Extrait les pièces avec leur bbox, sans remapping inutile"""
    result = pieces_pred[0]
    detections = []

    for box in result.boxes:
        # xywh en pixels de l'image originale
        x, y, w, h = box.xywh[0].cpu().numpy()
        label = result.names[int(box.cls[0])]
        detections.append({"label": label, "bbox": (int(x), int(y), int(w), int(h))})

    return detections

import numpy as np
import cv2

def pieces_to_board(detected_boxes, matrix, board_size=800):
    board_array = [[None for _ in range(8)] for _ in range(8)]

    for d in detected_boxes:
        x, y, w, h = d["bbox"]

        # Points multiples sur la pièce pour stabilité
        points = np.array([
            [x + w/2, y + h*0.2],   # haut
            [x + w/2, y + h/2],     # centre
            [x + w/2, y + h*0.8]    # bas
        ], dtype=np.float32).reshape(-1,1,2)

        # Transformation perspective
        warped_points = cv2.perspectiveTransform(points, matrix)
        wy_values = warped_points[:,0,1]  # coordonnées y après warp

        # Prendre le percentile haut (25%) pour éviter décalage
        wy_percentile = np.percentile(wy_values, 25)

        # Normaliser et calculer rank/file
        nx = np.clip(np.mean(warped_points[:,0,0]) / board_size, 0, 0.999)
        ny = np.clip(wy_percentile / board_size, 0, 0.999)

        file = min(max(int(nx * 8), 0), 7)
        rank = min(max(int(ny * 8), 0), 7)

        board_array[rank][file] = d["label"]

    return board_array


def board_to_fen(board):
    map_fen = {
        "w_pawn": "P", "w_knight": "N", "w_bishop": "B",
        "w_rook": "R", "w_queen": "Q", "w_king": "K",
        "b_pawn": "p", "b_knight": "n", "b_bishop": "b",
        "b_rook": "r", "b_queen": "q", "b_king": "k",
    }
    rows = []
    for rank in board:
        empty = 0
        row = ""
        for sq in rank:
            if sq is None:
                empty += 1
            else:
                if empty:
                    row += str(empty)
                    empty = 0
                row += map_fen[sq]
        if empty:
            row += str(empty)
        rows.append(row)
    return "/".join(rows)

if __name__ == "__main__":
    edges_detector = Detector("../assets/models/edges.pt")
    pieces_detector = Detector("../assets/models/unified-nano-refined.pt")
    #image_path = "./test/1.png"
    image_path = "../training/datasets/pieces/unified/test/images/659_jpg.rf.0009cadea8df487a76d6960a28b9d811.jpg"
    image = cv2.imread(image_path)

    edges_pred = edges_detector.make_prediction(image_path)
    pieces_pred = pieces_detector.make_prediction(image_path)

    remap_width = 800
    remap_height = 800

    board_manager = BoardManager(image)
    corners, matrix = board_manager.extract_corners(edges_pred[0], (remap_width, remap_height))

    detections = extract_pieces(pieces_pred)

    board = pieces_to_board(detections, matrix, remap_width)

    # FEN
    fen = board_to_fen(board)
    print("FEN:", fen)

    frame = pieces_pred[0].plot()
    cv2.namedWindow("Pred", cv2.WINDOW_NORMAL)
    cv2.imshow("Pred", frame)
    cv2.waitKey(0)
    cv2.destroyAllWindows()