"""
Анализ схемы нарезки спутниковых снимков в датасете UAV-GeoLoc.

Скрипт определяет crop_size, stride и overlap для каждой сцены,
сопоставляя кропы с исходным merge.tif через попиксельное сравнение.

Результат: stride = crop_size // 2 (50% overlap) для всех сцен.
Naming: crop_X_Y.png — X по ширине (col), Y по высоте (row).
Позиция в merge.tif: merge[Y*stride : Y*stride+crop_size, X*stride : X*stride+crop_size]
"""

import glob
import os

import numpy as np
from PIL import Image

Image.MAX_IMAGE_PIXELS = None  # некоторые merge.tif очень большие


def analyze_scene(scene_db_dir: str) -> dict:
    """Определяет параметры нарезки для одной сцены.

    Args:
        scene_db_dir: путь к папке DB сцены (содержит merge.tif и img/).

    Returns:
        dict с ключами: merge_size, crop_size, grid, stride, overlap.
    """
    merge_path = os.path.join(scene_db_dir, "merge.tif")
    img_dir = os.path.join(scene_db_dir, "img")

    merge = np.array(Image.open(merge_path))
    mh, mw = merge.shape[:2]

    # Размер кропа
    c00 = np.array(Image.open(os.path.join(img_dir, "crop_0_0.png")))
    ch, cw = c00.shape[:2]

    # Размер сетки
    crops = os.listdir(img_dir)
    xs, ys = [], []
    for name in crops:
        parts = name.replace("crop_", "").replace(".png", "").split("_")
        xs.append(int(parts[0]))
        ys.append(int(parts[1]))
    grid_x, grid_y = max(xs) + 1, max(ys) + 1

    # Проверяем что crop_0_0 начинается с (0, 0)
    assert np.array_equal(c00, merge[0:ch, 0:cw, :3]), "crop_0_0 не совпадает с merge[0:ch, 0:cw]"

    # Ищем stride по X: сдвигаем crop_1_0 вдоль ширины merge
    c10 = np.array(Image.open(os.path.join(img_dir, "crop_1_0.png")))
    stride_x = None
    for s in range(1, cw + 1):
        if s + cw <= mw and np.array_equal(c10, merge[0:ch, s:s + cw, :3]):
            stride_x = s
            break
    assert stride_x is not None, "Не удалось найти stride по X"

    # Ищем stride по Y: сдвигаем crop_0_1 вдоль высоты merge
    c01 = np.array(Image.open(os.path.join(img_dir, "crop_0_1.png")))
    stride_y = None
    for s in range(1, ch + 1):
        if s + ch <= mh and np.array_equal(c01, merge[s:s + ch, 0:cw, :3]):
            stride_y = s
            break
    assert stride_y is not None, "Не удалось найти stride по Y"

    return {
        "merge_size": (mw, mh),
        "crop_size": (cw, ch),
        "grid": (grid_x, grid_y),
        "stride": (stride_x, stride_y),
        "overlap": (cw - stride_x, ch - stride_y),
    }


def main():
    base = os.path.dirname(os.path.abspath(__file__))

    patterns = [
        os.path.join(base, "Country", "*", "*", "*", "DB"),
        os.path.join(base, "Terrain", "*", "*", "DB"),
    ]

    scene_dirs = []
    for pat in patterns:
        scene_dirs.extend(sorted(glob.glob(pat)))

    print(f"{'Scene':<50} {'merge WxH':>14} {'crop':>8} {'grid':>8} {'stride':>8} {'overlap':>8}")
    print("-" * 100)

    for scene_db in scene_dirs:
        if not os.path.isfile(os.path.join(scene_db, "merge.tif")):
            continue

        # Короткое имя сцены
        rel = os.path.relpath(scene_db, base)
        name = rel.replace("/DB", "")

        try:
            info = analyze_scene(scene_db)
            mw, mh = info["merge_size"]
            cw, ch = info["crop_size"]
            gx, gy = info["grid"]
            sx, sy = info["stride"]
            ox, oy = info["overlap"]
            print(
                f"{name:<50} {mw:>6}x{mh:<6} {cw:>3}x{ch:<4} {gx:>3}x{gy:<4} {sx:>3}x{sy:<4} {ox:>3}x{oy:<4}"
            )
        except Exception as e:
            print(f"{name:<50} ERROR: {e}")


if __name__ == "__main__":
    main()