"""
Lab012. Сравнение цветного и серого кадров
для обнаружения человека и препятствий.

Создаются четыре варианта:

1. Цветной 320x240, JPEG quality 30.
2. Серый 320x240, JPEG quality 30.
3. Цветной 320x240, JPEG quality 15.
4. Серый 320x240, JPEG quality 15.

Для каждого варианта рассчитываются:

- размер JPEG;
- количество фрагментов;
- полный объём DATA + ACK;
- время передачи;
- стоимость сохранения цвета.
"""

from math import ceil
from pathlib import Path

from PIL import (
    Image,
    ImageDraw,
    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/lab012_safety_scene.jpg"
)

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

FRAME_SIZE = (320, 240)

FRAGMENT_DATA_SIZE = 512

IMAGE_ID_BASE = 2026071500

BITRATES_KBPS = [
    10,
    20,
    50,
]


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

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

    Изображение:
    - автоматически поворачивается по EXIF;
    - переводится в RGB;
    - масштабируется с сохранением пропорций;
    - при необходимости дополняется чёрными полями.
    """

    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,
    output_path: Path,
    quality: int,
) -> None:
    """
    Сохранить изображение в JPEG.
    """

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

    image.save(
        output_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:
    """
    Рассчитать параметры передачи JPEG
    с использованием текущего протокола.
    """

    image_bytes = file_path.read_bytes()

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

    total_data_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,
        )

        total_data_bytes += len(packet)

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

    total_ack_bytes = (
        len(fragments)
        * len(ack_packet)
    )

    total_radio_bytes = (
        total_data_bytes
        + total_ack_bytes
    )

    transfer_times = {}

    for bitrate_kbps in BITRATES_KBPS:

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

    return {
        "file_size": len(image_bytes),
        "fragment_count": len(fragments),
        "total_data_bytes": total_data_bytes,
        "total_ack_bytes": total_ack_bytes,
        "total_radio_bytes": total_radio_bytes,
        "times": transfer_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
)


# ============================================================
# Создание вариантов
# ============================================================

variant_definitions = [
    {
        "name": "color_320_q30",
        "image": color_frame,
        "quality": 30,
        "path": (
            OUTPUT_DIRECTORY
            / "01_color_320_q30.jpg"
        ),
    },
    {
        "name": "gray_320_q30",
        "image": gray_frame,
        "quality": 30,
        "path": (
            OUTPUT_DIRECTORY
            / "02_gray_320_q30.jpg"
        ),
    },
    {
        "name": "color_320_q15",
        "image": color_frame,
        "quality": 15,
        "path": (
            OUTPUT_DIRECTORY
            / "03_color_320_q15.jpg"
        ),
    },
    {
        "name": "gray_320_q15",
        "image": gray_frame,
        "quality": 15,
        "path": (
            OUTPUT_DIRECTORY
            / "04_gray_320_q15.jpg"
        ),
    },
]


for variant in variant_definitions:

    save_jpeg(
        image=variant["image"],
        output_path=variant["path"],
        quality=variant["quality"],
    )


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

results = []

for variant_index, variant in enumerate(
    variant_definitions
):

    transfer_result = estimate_transfer(
        file_path=variant["path"],
        image_id=IMAGE_ID_BASE + variant_index,
    )

    results.append(
        {
            "name": variant["name"],
            "path": variant["path"],
            "quality": variant["quality"],
            **transfer_result,
        }
    )


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

columns = 2
cell_width = 440
cell_height = 340

rows = ceil(
    len(results)
    / columns
)

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

draw = ImageDraw.Draw(
    comparison_image
)


for result_index, result in enumerate(results):

    column = result_index % columns
    row = result_index // columns

    cell_x = column * cell_width
    cell_y = row * cell_height

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

        preview = image.convert("RGB")

        # NEAREST не сглаживает пиксели при увеличении.
        # Поэтому дефекты сжатия будут хорошо видны.
        preview = preview.resize(
            (400, 300),
            Image.Resampling.NEAREST,
        )

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

    paste_y = cell_y + 35

    comparison_image.paste(
        preview,
        (paste_x, paste_y),
    )

    label = (
        f"{result['name']}\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(
        (cell_x + 10, cell_y + 5),
        label,
        fill="black",
        spacing=3,
    )


comparison_path = (
    OUTPUT_DIRECTORY
    / "lab012_comparison.jpg"
)

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


# ============================================================
# Поиск результатов по имени
# ============================================================

results_by_name = {
    result["name"]: result
    for result in results
}

color_q30 = results_by_name[
    "color_320_q30"
]

gray_q30 = results_by_name[
    "gray_320_q30"
]

color_q15 = results_by_name[
    "color_320_q15"
]

gray_q15 = results_by_name[
    "gray_320_q15"
]


# ============================================================
# Расчёт цены сохранения цвета
# ============================================================

color_overhead_q30_bytes = (
    color_q30["file_size"]
    - gray_q30["file_size"]
)

color_overhead_q30_percent = (
    color_overhead_q30_bytes
    / gray_q30["file_size"]
    * 100
)

color_overhead_q15_bytes = (
    color_q15["file_size"]
    - gray_q15["file_size"]
)

color_overhead_q15_percent = (
    color_overhead_q15_bytes
    / gray_q15["file_size"]
    * 100
)


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

print(
    "=== Lab012. Цветной или серый кадр ==="
)

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

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


print("\nСравнение вариантов:")

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

print("-" * 83)


for result in results:

    print(
        f"{result['name']:<22}"
        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Цена сохранения цвета при quality 30:")

print(
    format_size(
        color_overhead_q30_bytes
    ),
    f"или {color_overhead_q30_percent:.1f} %",
)


print("\nЦена сохранения цвета при quality 15:")

print(
    format_size(
        color_overhead_q15_bytes
    ),
    f"или {color_overhead_q15_percent:.1f} %",
)


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

print(comparison_path)


# ============================================================
# Проверки
# ============================================================

assert len(results) == 4

assert comparison_path.exists()

assert all(
    result["fragment_count"] > 0
    for result in results
)

assert color_q30["file_size"] > 0
assert gray_q30["file_size"] > 0
assert color_q15["file_size"] > 0
assert gray_q15["file_size"] > 0


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