---

import cv2
import numpy as np
import onnxruntime as ort

CLASS_NAMES = [
    "board",        # 0
    "white_king",   # 1
    "white_queen",  # 2
    "white_rook",   # 3
    "white_bishop", # 4
    "white_knight", # 5
    "white_pawn",   # 6
    "black_king",   # 7
    "black_queen",  # 8
    "black_rook",   # 9
    "black_bishop", # 10
    "black_knight", # 11
    "black_pawn"    # 12
]

# Маппинг названий в стандартные символы FEN
FEN_MAP = {
    "white_king": "K",
    "white_queen": "Q",
    "white_rook": "R",
    "white_bishop": "B",
    "white_knight": "N",
    "white_pawn": "P",
    "black_king": "k",
    "black_queen": "q",
    "black_rook": "r",
    "black_bishop": "b",
    "black_knight": "n",
    "black_pawn": "p"
}

class ChessAnalyzer:
    def __init__(self, model_path="best.onnx", conf_threshold=0.35):
        self.conf_threshold = conf_threshold
        self.session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
        self.input_name = self.session.get_inputs()[0].name

    def preprocess(self, img):
        h, w = img.shape[:2]
        resized = cv2.resize(img, (640, 640), interpolation=cv2.INTER_AREA)
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        tensor = rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
        tensor = np.expand_dims(tensor, axis=0)
        return tensor, w, h

    def analyze(self, image_input):
        if isinstance(image_input, str):
            img = cv2.imread(image_input)
        else:
            img = image_input

        if img is None:
            print("Ошибка: Не удалось загрузить изображение.")
            return []

        tensor, orig_w, orig_h = self.preprocess(img)
        outputs = self.session.run(None, {self.input_name: tensor})[0]
        predictions = np.squeeze(outputs).T
        
        boxes, confidences, class_ids = [], [], []
        x_factor = orig_w / 640.0
        y_factor = orig_h / 640.0

        for pred in predictions:
            scores = pred[4:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]

            if confidence >= self.conf_threshold:
                cx, cy, w, h = pred[0], pred[1], pred[2], pred[3]
                x1 = int((cx - w / 2) * x_factor)
                y1 = int((cy - h / 2) * y_factor)
                x2 = int((cx + w / 2) * x_factor)
                y2 = int((cy + h / 2) * y_factor)

                boxes.append([x1, y1, x2 - x1, y2 - y1])
                confidences.append(float(confidence))
                class_ids.append(class_id)

        indices = cv2.dnn.NMSBoxes(boxes, confidences, self.conf_threshold, 0.45)
        
        results = []
        if len(indices) > 0:
            for i in indices.flatten():
                x, y, w, h = boxes[i]
                cid = class_ids[i]
                name = CLASS_NAMES[cid] if cid < len(CLASS_NAMES) else f"class_{cid}"
                if name != "board":  # Игнорируем класс самой доски
                    results.append({
                        "name": name,
                        "box": [x, y, x + w, y + h],
                        "confidence": confidences[i]
                    })

        # Генерируем FEN и записываем в human.txt
        fen = self.export_to_fen(results, orig_w, orig_h)
        self.save_human_txt(fen)

        return results

    def export_to_fen(self, detections, img_w, img_h):
        """Преобразует список детектированных фигур в FEN-строку"""
        grid = [[None for _ in range(8)] for _ in range(8)]
        cell_w = img_w / 8.0
        cell_h = img_h / 8.0

        for det in detections:
            x1, y1, x2, y2 = det["box"]
            
            # Центр фигуры по горизонтали, ближе к основанию по вертикали
            cx = (x1 + x2) / 2.0
            cy = y2 - (y2 - y1) * 0.2  # 80% высоты рамки (ближе к нижней части фигуры)

            col = int(cx // cell_w)
            row = int(cy // cell_h)

            # Ограничиваем индексы от 0 до 7
            col = max(0, min(7, col))
            row = max(0, min(7, row))

            piece_char = FEN_MAP.get(det["name"])
            if piece_char:
                # Если клетка уже занята, сохраняем фигуру с большей уверенностью
                grid[row][col] = piece_char

        # Сборка FEN
        fen_rows = []
        for row in grid:
            empty_count = 0
            row_str = ""
            for cell in row:
                if cell is None:
                    empty_count += 1
                else:
                    if empty_count > 0:
                        row_str += str(empty_count)
                        empty_count = 0
                    row_str += cell
            if empty_count > 0:
                row_str += str(empty_count)
            fen_rows.append(row_str)

        return "/".join(fen_rows)

    def save_human_txt(self, fen_string, filepath="human.txt"):
        """Сохраняет FEN в файл в формате [FEN]"""
        content = f"[{fen_string}]"
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(content)
        print(f"Сохранено в {filepath}: {content}")


# === Запуск ===
if __name__ == "__main__":
    analyzer = ChessAnalyzer("best.onnx")
    # Передаем вырезанную доску от crop.py
    analyzer.analyze("chess_board.jpg")