"""
Lab010. Оптимизация изображения для слабого радиоканала.

Программа формирует несколько вариантов одного кадра:

1. Исходный JPEG.
2. Цветной кадр до 640x480.
3. Серый кадр до 640x480.
4. Серый кадр до 320x240, JPEG quality 30.
5. Серый кадр до 320x240, JPEG quality 15.
6. Карта контуров 320x240.
7. Серый кадр с наложенными контурами.

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

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

from csv import DictWriter
from hashlib import sha256
from math import ceil
from pathlib import Path

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

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

FRAGMENT_DATA_SIZE = 512

IMAGE_ID_BASE = 2026071300

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


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

def resize_inside(
    image: Image.Image,
    maximum_size: tuple[int, int],
) -> Image.Image:
    """
    Уменьшить изображение с сохранением пропорций.

    Изображение не растягивается и не искажается.
    """

    resized_image = image.copy()

    resized_image.thumbnail(
        maximum_size,
        Image.Resampling.LANCZOS,
    )

    return resized_image


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_bytes(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_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,
        )

        total_data_packet_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_packet_bytes
        + total_ack_bytes
    )

    efficiency_percent = (
        len(image_bytes)
        / total_radio_bytes
        * 100
    )

    transfer_times = {}

    for bitrate_kbps in BITRATES_KBPS:

        bitrate_bits_per_second = (
            bitrate_kbps * 1000
        )

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

    return {
        "file_size_bytes": len(image_bytes),
        "fragment_count": len(fragments),
        "data_packet_bytes": total_data_packet_bytes,
        "ack_bytes": total_ack_bytes,
        "total_radio_bytes": total_radio_bytes,
        "efficiency_percent": efficiency_percent,
        "transfer_times": transfer_times,
        "sha256": sha256(image_bytes).hexdigest(),
    }


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

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:

    source_image.load()

    source_rgb = source_image.convert("RGB")

    original_width, original_height = (
        source_rgb.size
    )


# ============================================================
# Формирование вариантов изображения
# ============================================================

variant_paths = []

# ------------------------------------------------------------
# Вариант 0. Исходный JPEG без изменения
# ------------------------------------------------------------

variant_paths.append(
    (
        "00_original",
        SOURCE_PATH,
    )
)


# ------------------------------------------------------------
# Вариант 1. Цветной кадр до 640x480, quality 40
# ------------------------------------------------------------

color_640 = resize_inside(
    source_rgb,
    (640, 480),
)

color_640_path = (
    OUTPUT_DIRECTORY
    / "01_color_640_q40.jpg"
)

save_jpeg(
    color_640,
    color_640_path,
    quality=40,
)

variant_paths.append(
    (
        "01_color_640_q40",
        color_640_path,
    )
)


# ------------------------------------------------------------
# Вариант 2. Серый кадр до 640x480, quality 40
# ------------------------------------------------------------

gray_640 = ImageOps.grayscale(
    color_640
)

gray_640_path = (
    OUTPUT_DIRECTORY
    / "02_gray_640_q40.jpg"
)

save_jpeg(
    gray_640,
    gray_640_path,
    quality=40,
)

variant_paths.append(
    (
        "02_gray_640_q40",
        gray_640_path,
    )
)


# ------------------------------------------------------------
# Вариант 3. Серый кадр до 320x240, quality 30
# ------------------------------------------------------------

color_320 = resize_inside(
    source_rgb,
    (320, 240),
)

gray_320 = ImageOps.grayscale(
    color_320
)

gray_320_q30_path = (
    OUTPUT_DIRECTORY
    / "03_gray_320_q30.jpg"
)

save_jpeg(
    gray_320,
    gray_320_q30_path,
    quality=30,
)

variant_paths.append(
    (
        "03_gray_320_q30",
        gray_320_q30_path,
    )
)


# ------------------------------------------------------------
# Вариант 4. Серый кадр до 320x240, quality 15
# ------------------------------------------------------------

gray_320_q15_path = (
    OUTPUT_DIRECTORY
    / "04_gray_320_q15.jpg"
)

save_jpeg(
    gray_320,
    gray_320_q15_path,
    quality=15,
)

variant_paths.append(
    (
        "04_gray_320_q15",
        gray_320_q15_path,
    )
)


# ------------------------------------------------------------
# Вариант 5. Только контуры
# ------------------------------------------------------------

edge_map = gray_320.filter(
    ImageFilter.FIND_EDGES
)

edge_map = ImageOps.autocontrast(
    edge_map
)

# Инвертируем: белый фон, тёмные контуры.
edge_map = ImageOps.invert(
    edge_map
)

edges_path = (
    OUTPUT_DIRECTORY
    / "05_edges_320_q40.jpg"
)

save_jpeg(
    edge_map,
    edges_path,
    quality=40,
)

variant_paths.append(
    (
        "05_edges_320_q40",
        edges_path,
    )
)


# ------------------------------------------------------------
# Вариант 6. Серый кадр с усиленными контурами
# ------------------------------------------------------------

gray_edges_overlay = Image.blend(
    gray_320,
    edge_map,
    alpha=0.25,
)

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

save_jpeg(
    gray_edges_overlay,
    gray_edges_path,
    quality=30,
)

variant_paths.append(
    (
        "06_gray_edges_320_q30",
        gray_edges_path,
    )
)


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

results = []

for variant_number, (
    variant_name,
    variant_path,
) in enumerate(variant_paths):

    with Image.open(variant_path) as variant_image:

        width, height = variant_image.size
        mode = variant_image.mode

    transfer_result = estimate_transfer(
        file_path=variant_path,
        image_id=IMAGE_ID_BASE + variant_number,
    )

    result = {
        "name": variant_name,
        "path": str(variant_path),
        "width": width,
        "height": height,
        "mode": mode,
        **transfer_result,
    }

    results.append(result)


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

CONTACT_SHEET_COLUMNS = 3
CONTACT_SHEET_CELL_WIDTH = 360
CONTACT_SHEET_CELL_HEIGHT = 290

contact_sheet_rows = ceil(
    len(results)
    / CONTACT_SHEET_COLUMNS
)

contact_sheet = Image.new(
    "RGB",
    (
        CONTACT_SHEET_COLUMNS
        * CONTACT_SHEET_CELL_WIDTH,
        contact_sheet_rows
        * CONTACT_SHEET_CELL_HEIGHT,
    ),
    "white",
)

draw = ImageDraw.Draw(
    contact_sheet
)

for result_index, result in enumerate(results):

    column = (
        result_index
        % CONTACT_SHEET_COLUMNS
    )

    row = (
        result_index
        // CONTACT_SHEET_COLUMNS
    )

    cell_x = (
        column
        * CONTACT_SHEET_CELL_WIDTH
    )

    cell_y = (
        row
        * CONTACT_SHEET_CELL_HEIGHT
    )

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

        preview = variant_image.convert("RGB")

        preview.thumbnail(
            (330, 220),
            Image.Resampling.LANCZOS,
        )

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

        paste_y = (
            cell_y + 35
        )

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

    label = (
        f"{result['name']}\n"
        f"{format_bytes(result['file_size_bytes'])}, "
        f"{result['fragment_count']} fragments"
    )

    draw.text(
        (cell_x + 10, cell_y + 8),
        label,
        fill="black",
    )

contact_sheet_path = (
    OUTPUT_DIRECTORY
    / "lab010_comparison.jpg"
)

contact_sheet.save(
    contact_sheet_path,
    format="JPEG",
    quality=90,
)


# ============================================================
# Сохранение таблицы CSV
# ============================================================

csv_path = (
    OUTPUT_DIRECTORY
    / "lab010_results.csv"
)

csv_fieldnames = [
    "name",
    "path",
    "width",
    "height",
    "mode",
    "file_size_bytes",
    "fragment_count",
    "total_radio_bytes",
    "efficiency_percent",
    "time_5_kbps",
    "time_10_kbps",
    "time_20_kbps",
    "time_50_kbps",
    "time_100_kbps",
    "time_20_kbps_with_retries_x2",
]

with csv_path.open(
    "w",
    newline="",
    encoding="utf-8-sig",
) as csv_file:

    writer = DictWriter(
        csv_file,
        fieldnames=csv_fieldnames,
    )

    writer.writeheader()

    for result in results:

        writer.writerow(
            {
                "name": result["name"],
                "path": result["path"],
                "width": result["width"],
                "height": result["height"],
                "mode": result["mode"],
                "file_size_bytes": (
                    result["file_size_bytes"]
                ),
                "fragment_count": (
                    result["fragment_count"]
                ),
                "total_radio_bytes": (
                    result["total_radio_bytes"]
                ),
                "efficiency_percent": (
                    f"{result['efficiency_percent']:.2f}"
                ),
                "time_5_kbps": (
                    result["transfer_times"][5]
                ),
                "time_10_kbps": (
                    result["transfer_times"][10]
                ),
                "time_20_kbps": (
                    result["transfer_times"][20]
                ),
                "time_50_kbps": (
                    result["transfer_times"][50]
                ),
                "time_100_kbps": (
                    result["transfer_times"][100]
                ),
                "time_20_kbps_with_retries_x2": (
                    result["transfer_times"][20]
                    * 2
                ),
            }
        )


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

print(
    "=== Lab010. Оптимизация изображения ==="
)

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

print(
    original_width,
    "x",
    original_height,
)


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

print(
    f"{'Вариант':<28}"
    f"{'Размер':>12}"
    f"{'Фрагм.':>9}"
    f"{'10 кбит/с':>13}"
    f"{'20 кбит/с':>13}"
    f"{'50 кбит/с':>13}"
    f"{'20 кбит/с x2':>16}"
)

print("-" * 104)

for result in results:

    print(
        f"{result['name']:<28}"
        f"{format_bytes(result['file_size_bytes']):>12}"
        f"{result['fragment_count']:>9}"
        f"{format_duration(result['transfer_times'][10]):>13}"
        f"{format_duration(result['transfer_times'][20]):>13}"
        f"{format_duration(result['transfer_times'][50]):>13}"
        f"{format_duration(result['transfer_times'][20] * 2):>16}"
    )


# ============================================================
# Итоговые кандидаты
# ============================================================

smallest_result = min(
    results,
    key=lambda item: item["file_size_bytes"],
)

operator_candidate = next(
    result
    for result in results
    if result["name"] == "03_gray_320_q30"
)


print("\nСамый маленький файл:")

print(
    smallest_result["name"],
    "-",
    format_bytes(
        smallest_result["file_size_bytes"]
    ),
)


print("\nБазовый кандидат для оператора:")

print(
    operator_candidate["name"],
    "-",
    format_bytes(
        operator_candidate["file_size_bytes"]
    ),
    "-",
    format_duration(
        operator_candidate[
            "transfer_times"
        ][20]
    ),
    "при 20 кбит/с без повторов",
)


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

print(contact_sheet_path)


print("\nТаблица результатов CSV:")

print(csv_path)


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

assert len(results) == 7

assert contact_sheet_path.exists()

assert csv_path.exists()

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

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