"""
Lab011. Проверка облегчённых режимов изображения
на сцене, похожей на вид с камеры ровера.

Создаются варианты:
1. Grayscale 320x240, JPEG quality 30.
2. Grayscale 320x240, JPEG quality 15.
3. Grayscale 320x240 с наложением контуров, quality 30.

Для каждого варианта рассчитываются:
- размер JPEG;
- количество радиопакетов;
- время передачи при разных скоростях.
"""

from math import ceil
from pathlib import Path

from PIL import (
    Image,
    ImageChops,
    ImageDraw,
    ImageFilter,
    ImageOps,
)

from protocol.image_fragments import (
    encode_image_fragment,
    split_image_bytes,
)

from protocol.packet import (
    MESSAGE_TYPE_ACK,
    MESSAGE_TYPE_IMAGE_FRAGMENT,
    build_packet,
)


# ============================================================
# Настройки
# ============================================================

SOURCE_PATH = Path(
    "data/raw/lab011_scene.jpg"
)

OUTPUT_DIRECTORY = Path(
    "data/processed/lab011"
)

FRAME_SIZE = (320, 240)

FRAGMENT_DATA_SIZE = 512

IMAGE_ID_BASE = 2026071400

BITRATES_KBPS = [
    5,
    10,
    20,
    50,
]


# ============================================================
# Вспомогательные функции
# ============================================================

def prepare_frame(
    image: Image.Image,
) -> Image.Image:
    """
    Привести изображение к фиксированному кадру 320x240.

    Пропорции сохраняются.
    Недостающие области заполняются чёрным.
    """

    image = ImageOps.exif_transpose(image)
    image = image.convert("RGB")

    return ImageOps.pad(
        image,
        FRAME_SIZE,
        method=Image.Resampling.LANCZOS,
        color="black",
        centering=(0.5, 0.5),
    )


def save_jpeg(
    image: Image.Image,
    path: Path,
    quality: int,
) -> None:
    """
    Сохранить изображение в JPEG.
    """

    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    image.save(
        path,
        format="JPEG",
        quality=quality,
        optimize=True,
    )


def format_size(
    byte_count: int,
) -> str:
    """
    Представить размер в КиБ.
    """

    return f"{byte_count / 1024:.2f} КиБ"


def format_duration(
    seconds: float,
) -> str:
    """
    Представить длительность передачи.
    """

    if seconds < 1:
        return f"{seconds * 1000:.0f} мс"

    if seconds < 60:
        return f"{seconds:.2f} с"

    minutes = int(seconds // 60)
    remaining_seconds = seconds % 60

    return (
        f"{minutes} мин "
        f"{remaining_seconds:.1f} с"
    )


def estimate_transfer(
    file_path: Path,
    image_id: int,
) -> dict:
    """
    Рассчитать полный объём DATA + ACK.
    """

    image_bytes = file_path.read_bytes()

    fragments = split_image_bytes(
        image_bytes=image_bytes,
        image_id=image_id,
        fragment_data_size=FRAGMENT_DATA_SIZE,
    )

    data_packet_bytes = 0

    for fragment in fragments:

        fragment_payload = encode_image_fragment(
            fragment
        )

        packet = build_packet(
            payload=fragment_payload,
            message_type=MESSAGE_TYPE_IMAGE_FRAGMENT,
            sequence_number=fragment.fragment_index,
        )

        data_packet_bytes += len(packet)

    ack_packet = build_packet(
        payload=b"",
        message_type=MESSAGE_TYPE_ACK,
        sequence_number=0,
    )

    ack_bytes = (
        len(fragments)
        * len(ack_packet)
    )

    total_radio_bytes = (
        data_packet_bytes
        + ack_bytes
    )

    times = {}

    for bitrate_kbps in BITRATES_KBPS:

        times[bitrate_kbps] = (
            total_radio_bytes
            * 8
            / (bitrate_kbps * 1000)
        )

    return {
        "file_size": len(image_bytes),
        "fragment_count": len(fragments),
        "total_radio_bytes": total_radio_bytes,
        "times": times,
    }


# ============================================================
# Проверка исходного файла
# ============================================================

if not SOURCE_PATH.exists():
    raise FileNotFoundError(
        f"Не найден файл: {SOURCE_PATH}"
    )

OUTPUT_DIRECTORY.mkdir(
    parents=True,
    exist_ok=True,
)


# ============================================================
# Подготовка базового кадра
# ============================================================

with Image.open(SOURCE_PATH) as source_image:

    original_size = source_image.size

    color_frame = prepare_frame(
        source_image
    )

gray_frame = ImageOps.grayscale(
    color_frame
)


# ============================================================
# Вариант 1. Grayscale, quality 30
# ============================================================

gray_q30_path = (
    OUTPUT_DIRECTORY
    / "01_gray_320_q30.jpg"
)

save_jpeg(
    gray_frame,
    gray_q30_path,
    quality=30,
)


# ============================================================
# Вариант 2. Grayscale, quality 15
# ============================================================

gray_q15_path = (
    OUTPUT_DIRECTORY
    / "02_gray_320_q15.jpg"
)

save_jpeg(
    gray_frame,
    gray_q15_path,
    quality=15,
)


# ============================================================
# Вариант 3. Grayscale + контуры
# ============================================================

edge_map = gray_frame.filter(
    ImageFilter.FIND_EDGES
)

edge_map = ImageOps.autocontrast(
    edge_map
)

# Оставляем преимущественно сильные контуры.
binary_edges = edge_map.point(
    lambda value: 255 if value >= 45 else 0
)

# После инверсии контуры становятся чёрными,
# а фон — белым.
dark_edges = ImageOps.invert(
    binary_edges
)

# Сохраняем исходный серый фон и добавляем
# поверх него тёмные линии контуров.
gray_with_edges = ImageChops.darker(
    gray_frame,
    dark_edges
)

gray_edges_path = (
    OUTPUT_DIRECTORY
    / "03_gray_edges_320_q30.jpg"
)

save_jpeg(
    gray_with_edges,
    gray_edges_path,
    quality=30,
)


# ============================================================
# Список вариантов
# ============================================================

variants = [
    (
        "gray_320_q30",
        gray_q30_path,
    ),
    (
        "gray_320_q15",
        gray_q15_path,
    ),
    (
        "gray_edges_320_q30",
        gray_edges_path,
    ),
]


# ============================================================
# Расчёт параметров
# ============================================================

results = []

for variant_index, (
    variant_name,
    variant_path,
) in enumerate(variants):

    result = estimate_transfer(
        file_path=variant_path,
        image_id=IMAGE_ID_BASE + variant_index,
    )

    results.append(
        {
            "name": variant_name,
            "path": variant_path,
            **result,
        }
    )


# ============================================================
# Создание сравнительной картинки
# ============================================================

preview_variants = [
    (
        "SOURCE PREVIEW",
        color_frame,
        None,
    )
]

for result in results:

    with Image.open(result["path"]) as image:

        preview_variants.append(
            (
                result["name"],
                image.convert("RGB").copy(),
                result,
            )
        )


columns = 2
cell_width = 440
cell_height = 340

rows = ceil(
    len(preview_variants)
    / columns
)

comparison_image = Image.new(
    "RGB",
    (
        columns * cell_width,
        rows * cell_height,
    ),
    "white",
)

draw = ImageDraw.Draw(
    comparison_image
)

for index, (
    label,
    preview,
    result,
) in enumerate(preview_variants):

    column = index % columns
    row = index // columns

    x = column * cell_width
    y = row * cell_height

    preview = preview.copy()

    preview.thumbnail(
        (400, 270),
        Image.Resampling.NEAREST,
    )

    paste_x = (
        x
        + (cell_width - preview.width) // 2
    )

    comparison_image.paste(
        preview,
        (paste_x, y + 55),
    )

    if result is None:

        text = (
            f"{label}\n"
            f"original: "
            f"{original_size[0]}x{original_size[1]}"
        )

    else:

        text = (
            f"{label}\n"
            f"{format_size(result['file_size'])}, "
            f"{result['fragment_count']} fragments, "
            f"{format_duration(result['times'][20])} "
            f"at 20 kbps"
        )

    draw.multiline_text(
        (x + 10, y + 10),
        text,
        fill="black",
        spacing=4,
    )


comparison_path = (
    OUTPUT_DIRECTORY
    / "lab011_comparison.jpg"
)

comparison_image.save(
    comparison_path,
    format="JPEG",
    quality=90,
)


# ============================================================
# Вывод результатов
# ============================================================

print(
    "=== Lab011. Режимы кадра для ровера ==="
)

print("\nИсходное разрешение:")

print(
    original_size[0],
    "x",
    original_size[1],
)

print("\nРезультаты:")

print(
    f"{'Вариант':<27}"
    f"{'Размер':>12}"
    f"{'Пакеты':>10}"
    f"{'10 кбит/с':>13}"
    f"{'20 кбит/с':>13}"
    f"{'50 кбит/с':>13}"
)

print("-" * 88)

for result in results:

    print(
        f"{result['name']:<27}"
        f"{format_size(result['file_size']):>12}"
        f"{result['fragment_count']:>10}"
        f"{format_duration(result['times'][10]):>13}"
        f"{format_duration(result['times'][20]):>13}"
        f"{format_duration(result['times'][50]):>13}"
    )


print("\nСравнительная картинка:")

print(comparison_path)


print(
    "\nПроверка пройдена: "
    "режимы кадра созданы и рассчитаны."
)