---

import numpy as np
from PIL import Image

def kmeans_compress(image_path, output_path, k=8, max_iters=20, tol=1e-2):
    # 1. Загрузка изображения
    img = Image.open(image_path).convert('RGB')
    img_np = np.array(img, dtype=np.float32)
    height, width, channels = img_np.shape
    
    # Сохраняем исходную матрицу в формате uint8 для сборки финального коллажа
    original_uint8 = np.array(img, dtype=np.uint8)

    # Вытягиваем в плоскую таблицу пикселей
    pixels = img_np.reshape(-1, channels)

    # 2. Инициализация случайных центроидов
    initial_indices = np.random.choice(pixels.shape[0], k, replace=False)
    centroids = pixels[initial_indices]

    print(f"Начало кластеризации для K={k}...")

    # Главный цикл K-means
    for i in range(max_iters):
        # Расчет расстояний и назначение кластеров
        distances = np.linalg.norm(pixels[:, np.newaxis, :] - centroids, axis=2)
        labels = np.argmin(distances, axis=1)

        # Пересчет центроидов
        new_centroids = np.zeros_like(centroids)
        for c in range(k):
            cluster_pixels = pixels[labels == c]
            if len(cluster_pixels) > 0:
                new_centroids[c] = cluster_pixels.mean(axis=0)
            else:
                new_centroids[c] = pixels[np.random.choice(pixels.shape[0])]

        # Проверка сходимости
        center_shift = np.linalg.norm(centroids - new_centroids)
        if center_shift < tol:
            centroids = new_centroids
            print(f"Алгоритм сошелся на итерации {i+1}")
            break
            
        centroids = new_centroids
        print(f"Итерация {i+1}/{max_iters}, сдвиг: {center_shift:.4f}")

    # --- ВЫВОД КЛАСТЕРОВ В КОНСОЛЬ ---
    print("\n" + "="*40)
    print(f" НАЙДЕННЫЕ КЛАСТЕРЫ (ЦВЕРА ПАЛИТРЫ) ДЛЯ K={k}:")
    print("="*40)
    for idx, color in enumerate(centroids):
        r, g, b = int(np.clip(color[0], 0, 255)), int(np.clip(color[1], 0, 255)), int(np.clip(color[2], 0, 255))
        print(f"Кластер {idx + 1}: RGB({r}, {g}, {b}) -> Порядковый номер в палитре")
    print("="*40 + "\n")

    # 3. Создание сжатого изображения
    compressed_pixels = centroids[labels]
    compressed_img_np = compressed_pixels.reshape(height, width, channels)
    compressed_img_np = np.clip(compressed_img_np, 0, 255).astype(np.uint8)

    # Сохраняем просто сжатую картинку
    Image.fromarray(compressed_img_np).save(output_path)
    print(f"Одиночное сжатое изображение сохранено в: {output_path}")

    # --- СБОРКА КОЛЛАЖА КАК В МЕТОДИЧКЕ (Рисунок 12) ---
    # Правая часть: оригинал сверху, сжатая снизу
    right_side = np.vstack([original_uint8, compressed_img_np])

    # Левая часть: вертикальные полосы кластеров
    palette_height = height * 2
    palette_width = width
    palette_np = np.zeros((palette_height, palette_width, 3), dtype=np.uint8)
    
    # Считаем ширину одной полосы цвета
    bar_width = palette_width // k
    for c in range(k):
        start_x = c * bar_width
        # Последняя полоса забирает остаток пикселей из-за округления
        end_x = (c + 1) * bar_width if c < k - 1 else palette_width
        
        # Закрашиваем полосу цветом текущего центроида
        r, g, b = int(np.clip(centroids[c][0], 0, 255)), int(np.clip(centroids[c][1], 0, 255)), int(np.clip(centroids[c][2], 0, 255))
        palette_np[:, start_x:end_x, :] = [r, g, b]

    # Соединяем левую палитру и правую стопку картинок по горизонтали
    final_report_image = np.hstack([palette_np, right_side])

    # Сохраняем готовый коллаж для отчета
    report_output_path = "report_result.jpg"
    Image.fromarray(final_report_image).save(report_output_path)
    print(f"Готовый коллаж для вставки в отчет сохранен в: {report_output_path}")

if __name__ == "__main__":
    user_input = input("Введите желаемое количество цветов (кластеров) K: ")
    try:
        k_value = int(user_input)
        if k_value <= 0:
            print("Ошибка: Количество кластеров должно быть больше нуля!")
        else:
            # Запускаем (убедись, что файл input.jpg лежит в той же папке)
            kmeans_compress("input.jpg", "output_python.jpg", k=k_value)
    except ValueError:
        print("Ошибка: Вы ввели не число!").