import cv2
import numpy as np
import random
import json

# Загрузка изображения
image = cv2.imread("a4keyboard_scan.jpg")
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# Применяем адаптивный порог
thresh = cv2.adaptiveThreshold(
    gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2
)

# Находим контуры и собираем размеры прямоугольников
rectangles = [
    cv2.boundingRect(contour)
    for contour in cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[
        0
    ]
    if (cv2.boundingRect(contour)[2] > 20 and cv2.boundingRect(contour)[3] > 20)
]

# Вычисляем медианные размеры
widths, heights = zip(*[(w, h) for _, _, w, h in rectangles])
median_width, median_height = np.median(widths), np.median(heights)

# Фильтрация прямоугольников по размеру
filtered_rectangles = [
    rect
    for rect in rectangles
    if 0.5 * median_width <= rect[2] <= 10 * median_width
    and 0.5 * median_height <= rect[3] <= 10 * median_height
]

# Порог для тёмного фона
darkness_threshold = 80
filtered_rectangles_dark = [
    rect
    for rect in filtered_rectangles
    if cv2.mean(gray[rect[1] : rect[1] + rect[3], rect[0] : rect[0] + rect[2]])[0]
    < darkness_threshold
]

# Уменьшаем размеры и закрашиваем области
padding = 8
for x, y, w, h in filtered_rectangles_dark:
    x_padded, y_padded = x + padding, y + padding
    w_padded, h_padded = max(w - 2 * padding, 1), max(h - 2 * padding, 1)
    roi = image[y_padded : y_padded + h_padded, x_padded : x_padded + w_padded]
    mean_color = cv2.mean(roi)[:3]

    # Генерация области с шумом
    noise = np.random.uniform(-10, 10, roi.shape).astype(np.int16)
    noisy_color = np.clip(np.array(mean_color, dtype=np.int16) + noise, 0, 255).astype(
        np.uint8
    )
    image[y_padded : y_padded + h_padded, x_padded : x_padded + w_padded] = noisy_color

# Масштабирование изображения и прямоугольников
scale_factor = 0.65
# image = cv2.resize(image, (image.shape[1] // 2, image.shape[0] // 2))
image = cv2.resize(
    image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR
)
scaled_rectangles = [
    (
        int(x * scale_factor),
        int(y * scale_factor),
        int(w * scale_factor),
        int(h * scale_factor),
    )
    for x, y, w, h in filtered_rectangles_dark
]

# Сортируем прямоугольники построчно
row_tolerance = 10
scaled_rectangles.sort(key=lambda rect: rect[1])

rows = []
current_row = [scaled_rectangles[0]]

for rect in scaled_rectangles[1:]:
    if abs(rect[1] - current_row[-1][1]) <= row_tolerance:
        current_row.append(rect)
    else:
        rows.append(current_row)
        current_row = [rect]

rows.append(current_row)

sorted_rectangles = [
    rect for row in rows for rect in sorted(row, key=lambda rect: rect[0])
]

# Сохраняем изображение и данные прямоугольников
cv2.imwrite("background_image.png", image)

with open("rectangles.json", "w") as f:
    json.dump(sorted_rectangles, f)