"""
Lab021. Оценка и компенсация частотного рассогласования BPSK.

Программа:

1. Загружает IQ-радиокадр из Lab018.
2. Добавляет постоянный фазовый поворот.
3. Добавляет заданное частотное рассогласование.
4. Добавляет неизвестную задержку.
5. Выполняет согласованную RRC-фильтрацию.
6. Ищет PREAMBLE + RADIO SYNC.
7. По известному маркеру оценивает:
   - начало кадра;
   - фазу символьной дискретизации;
   - начальный фазовый поворот;
   - частотную ошибку.
8. Сравнивает приём:
   - только с постоянной фазовой коррекцией;
   - с фазовой и частотной коррекцией.
9. Восстанавливает внутренний пакет и проверяет CRC.

Шум в этой лабораторной не добавляется.
Изучается только влияние частотной ошибки.
"""

from pathlib import Path
import struct

import matplotlib.pyplot as plt
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view
from scipy.signal import fftconvolve

from protocol.packet import (
    CRCError,
    MESSAGE_TYPE_TEXT,
    PacketError,
    parse_packet,
)


# ============================================================
# Параметры радиокадра
# ============================================================

RADIO_SYNC_WORD = 0xD391

PREAMBLE_BIT_COUNT = 64

SAMPLES_PER_SYMBOL = 32

SYMBOL_RATE = 20_000

SAMPLE_RATE = (
    SYMBOL_RATE
    * SAMPLES_PER_SYMBOL
)

RRC_ROLLOFF = 0.35
RRC_SPAN_SYMBOLS = 10


# ============================================================
# Контрольные данные
# ============================================================

EXPECTED_MESSAGE = "ПРИВЕТ SDR"
EXPECTED_SEQUENCE_NUMBER = 18


# ============================================================
# Искусственные искажения
# ============================================================

TEST_SAMPLE_DELAY = 11

TEST_PHASE_OFFSET_DEGREES = 37.0

FREQUENCY_OFFSETS_HZ = [
    0.0,
    10.0,
    25.0,
    50.0,
    100.0,
    250.0,
    500.0,
    1000.0,
]


# ============================================================
# Пути
# ============================================================

INPUT_IQ_PATH = Path(
    "data/processed/lab018/"
    "lab018_bpsk_tx_iq.npy"
)

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

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

GRAPH_PATH = (
    OUTPUT_DIRECTORY
    / "lab021_frequency_offset.png"
)

REPORT_PATH = (
    OUTPUT_DIRECTORY
    / "lab021_frequency_offset_report.txt"
)


# ============================================================
# Статусы приёма
# ============================================================

STATUS_SUCCESS = "SUCCESS"
STATUS_HEADER_ERROR = "HEADER ERROR"
STATUS_CRC_ERROR = "CRC ERROR"
STATUS_PACKET_ERROR = "PACKET ERROR"


# ============================================================
# Bytes и bits
# ============================================================

def bytes_to_bits(
    data: bytes,
) -> np.ndarray:
    """
    Преобразовать bytes в одномерный массив битов.
    """

    if not isinstance(data, bytes):
        raise TypeError(
            "data должен иметь тип bytes"
        )

    return np.unpackbits(
        np.frombuffer(
            data,
            dtype=np.uint8,
        )
    )


def bits_to_bytes(
    bits: np.ndarray,
) -> bytes:
    """
    Упаковать отдельные биты обратно в bytes.
    """

    bits = np.asarray(
        bits,
        dtype=np.uint8,
    )

    if bits.ndim != 1:
        raise ValueError(
            "bits должен быть одномерным массивом"
        )

    if len(bits) % 8 != 0:
        raise ValueError(
            "Количество битов должно быть кратно восьми"
        )

    if not np.all(
        (bits == 0) | (bits == 1)
    ):
        raise ValueError(
            "bits должен содержать только 0 и 1"
        )

    return np.packbits(
        bits
    ).tobytes()


# ============================================================
# BPSK
# ============================================================

def bpsk_modulate(
    bits: np.ndarray,
) -> np.ndarray:
    """
    Преобразовать биты в BPSK-символы.

        0 → -1
        1 → +1
    """

    bits = np.asarray(
        bits,
        dtype=np.uint8,
    )

    symbols = (
        2.0
        * bits.astype(np.float64)
        - 1.0
    )

    return symbols.astype(
        np.complex128
    )


def bpsk_demodulate(
    symbols: np.ndarray,
) -> np.ndarray:
    """
    Демодулировать BPSK по знаку компоненты I.
    """

    return (
        symbols.real >= 0.0
    ).astype(np.uint8)


# ============================================================
# Root Raised Cosine
# ============================================================

def root_raised_cosine_taps(
    rolloff: float,
    samples_per_symbol: int,
    span_symbols: int,
) -> np.ndarray:
    """
    Рассчитать коэффициенты RRC-фильтра.
    """

    if not 0.0 < rolloff <= 1.0:
        raise ValueError(
            "rolloff должен находиться в диапазоне 0...1"
        )

    if samples_per_symbol <= 0:
        raise ValueError(
            "samples_per_symbol должен быть положительным"
        )

    if span_symbols <= 0:
        raise ValueError(
            "span_symbols должен быть положительным"
        )

    if span_symbols % 2 != 0:
        raise ValueError(
            "span_symbols должен быть чётным"
        )

    half_sample_count = (
        span_symbols
        * samples_per_symbol
        // 2
    )

    sample_indexes = np.arange(
        -half_sample_count,
        half_sample_count + 1,
        dtype=np.float64,
    )

    time_values = (
        sample_indexes
        / samples_per_symbol
    )

    taps = np.zeros_like(
        time_values
    )

    beta = rolloff

    for index, time_value in enumerate(
        time_values
    ):

        if np.isclose(
            time_value,
            0.0,
        ):
            taps[index] = (
                1.0
                - beta
                + 4.0 * beta / np.pi
            )

            continue

        if np.isclose(
            abs(time_value),
            1.0 / (4.0 * beta),
        ):
            taps[index] = (
                beta
                / np.sqrt(2.0)
                * (
                    (
                        1.0
                        + 2.0 / np.pi
                    )
                    * np.sin(
                        np.pi
                        / (4.0 * beta)
                    )
                    + (
                        1.0
                        - 2.0 / np.pi
                    )
                    * np.cos(
                        np.pi
                        / (4.0 * beta)
                    )
                )
            )

            continue

        numerator = (
            np.sin(
                np.pi
                * time_value
                * (1.0 - beta)
            )
            + (
                4.0
                * beta
                * time_value
                * np.cos(
                    np.pi
                    * time_value
                    * (1.0 + beta)
                )
            )
        )

        denominator = (
            np.pi
            * time_value
            * (
                1.0
                - (
                    4.0
                    * beta
                    * time_value
                ) ** 2
            )
        )

        taps[index] = (
            numerator
            / denominator
        )

    taps /= np.sqrt(
        np.sum(
            taps ** 2
        )
    )

    return taps


# ============================================================
# Маркер радиокадра
# ============================================================

def build_frame_marker(
) -> tuple[np.ndarray, np.ndarray]:
    """
    Сформировать PREAMBLE + RADIO SYNC.
    """

    preamble_bits = np.tile(
        np.array(
            [1, 0],
            dtype=np.uint8,
        ),
        PREAMBLE_BIT_COUNT // 2,
    )

    sync_bits = bytes_to_bits(
        struct.pack(
            ">H",
            RADIO_SYNC_WORD,
        )
    )

    marker_bits = np.concatenate(
        [
            preamble_bits,
            sync_bits,
        ]
    )

    marker_symbols = bpsk_modulate(
        marker_bits
    )

    return marker_bits, marker_symbols


# ============================================================
# Добавление частотного смещения
# ============================================================

def apply_frequency_offset(
    iq_samples: np.ndarray,
    frequency_offset_hz: float,
    sample_rate: float,
    phase_offset_degrees: float,
) -> np.ndarray:
    """
    Добавить постоянную фазу и частотное рассогласование.

    В каждом следующем сэмпле фаза увеличивается на:

        2π · Δf / Fs
    """

    sample_indexes = np.arange(
        len(iq_samples),
        dtype=np.float64,
    )

    initial_phase_radians = np.deg2rad(
        phase_offset_degrees
    )

    phase_values = (
        initial_phase_radians
        + 2.0
        * np.pi
        * frequency_offset_hz
        * sample_indexes
        / sample_rate
    )

    return (
        iq_samples
        * np.exp(
            1j * phase_values
        )
    )


# ============================================================
# Поиск маркера и оценка CFO
# ============================================================

def find_frame_and_frequency_offset(
    matched_iq: np.ndarray,
    marker_symbols: np.ndarray,
    samples_per_symbol: int,
    symbol_rate: float,
) -> dict:
    """
    Найти радиокадр и оценить частотную ошибку.

    Для каждого возможного положения маркера:

    1. Умножаем принятые символы на известные
       символы маркера.

       Это удаляет BPSK-модуляцию:

           received × marker
           ≈ exp(j · phase)

    2. Измеряем среднее изменение фазы
       между соседними символами.

    3. Компенсируем полученный наклон фазы.

    4. Рассчитываем когерентную корреляцию.
    """

    marker_length = len(
        marker_symbols
    )

    marker_energy = float(
        np.sum(
            np.abs(marker_symbols) ** 2
        )
    )

    marker_indexes = np.arange(
        marker_length,
        dtype=np.float64,
    )

    best_result = None

    for sample_phase in range(
        samples_per_symbol
    ):

        symbol_samples = matched_iq[
            sample_phase::samples_per_symbol
        ]

        if len(symbol_samples) < marker_length:
            continue

        windows = sliding_window_view(
            symbol_samples,
            marker_length,
        )

        # Удаление известных знаков BPSK-маркера.
        despread_windows = (
            windows
            * marker_symbols[
                np.newaxis,
                :
            ]
        )

        # Изменение фазы между соседними символами.
        adjacent_products = (
            despread_windows[:, 1:]
            * np.conj(
                despread_windows[:, :-1]
            )
        )

        phase_increments = np.angle(
            np.sum(
                adjacent_products,
                axis=1,
            )
        )

        # Компенсация предполагаемой частотной ошибки.
        frequency_compensation = np.exp(
            -1j
            * phase_increments[:, np.newaxis]
            * marker_indexes[np.newaxis, :]
        )

        compensated_windows = (
            despread_windows
            * frequency_compensation
        )

        coherent_sums = np.sum(
            compensated_windows,
            axis=1,
        )

        window_energy = np.sum(
            np.abs(windows) ** 2,
            axis=1,
        )

        normalized_scores = (
            np.abs(coherent_sums)
            / (
                np.sqrt(
                    window_energy
                    * marker_energy
                )
                + 1e-12
            )
        )

        best_start_for_phase = int(
            np.argmax(
                normalized_scores
            )
        )

        best_score_for_phase = float(
            normalized_scores[
                best_start_for_phase
            ]
        )

        if (
            best_result is None
            or best_score_for_phase
            > best_result["score"]
        ):

            selected_phase_increment = float(
                phase_increments[
                    best_start_for_phase
                ]
            )

            selected_coherent_sum = (
                coherent_sums[
                    best_start_for_phase
                ]
            )

            estimated_frequency_offset_hz = (
                selected_phase_increment
                * symbol_rate
                / (2.0 * np.pi)
            )

            estimated_initial_phase = np.angle(
                selected_coherent_sum
            )

            best_result = {
                "score": best_score_for_phase,
                "sample_phase": sample_phase,
                "start_symbol_index": (
                    best_start_for_phase
                ),
                "symbol_samples": symbol_samples,
                "phase_increment": (
                    selected_phase_increment
                ),
                "frequency_offset_hz": (
                    estimated_frequency_offset_hz
                ),
                "initial_phase": (
                    estimated_initial_phase
                ),
            }

    if best_result is None:
        raise RuntimeError(
            "Не удалось обнаружить радиокадр"
        )

    return best_result


# ============================================================
# Восстановление внутреннего пакета
# ============================================================

def decode_radio_frame(
    corrected_symbols: np.ndarray,
    frame_start_symbol: int,
    marker_bits: np.ndarray,
) -> dict:
    """
    Демодулировать радиокадр и проверить внутренний пакет.
    """

    received_bits = bpsk_demodulate(
        corrected_symbols
    )

    available_bits = received_bits[
        frame_start_symbol:
    ]

    marker_bit_count = len(
        marker_bits
    )

    if len(available_bits) < marker_bit_count:
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": None,
            "message": None,
        }

    received_marker_bits = available_bits[
        :marker_bit_count
    ]

    marker_errors = int(
        np.count_nonzero(
            received_marker_bits
            != marker_bits
        )
    )

    radio_header_start = (
        PREAMBLE_BIT_COUNT
    )

    radio_header_end = (
        radio_header_start + 32
    )

    if len(available_bits) < radio_header_end:
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    try:
        radio_header = bits_to_bytes(
            available_bits[
                radio_header_start:
                radio_header_end
            ]
        )

        (
            received_sync,
            protocol_packet_length,
        ) = struct.unpack(
            ">HH",
            radio_header,
        )

    except (ValueError, struct.error):
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    if received_sync != RADIO_SYNC_WORD:
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    if not 1 <= protocol_packet_length <= 4096:
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    protocol_packet_start = (
        radio_header_end
    )

    protocol_packet_end = (
        protocol_packet_start
        + protocol_packet_length * 8
    )

    if len(available_bits) < protocol_packet_end:
        return {
            "status": STATUS_HEADER_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    try:
        protocol_packet = bits_to_bytes(
            available_bits[
                protocol_packet_start:
                protocol_packet_end
            ]
        )

        parsed_packet = parse_packet(
            protocol_packet
        )

    except CRCError:
        return {
            "status": STATUS_CRC_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    except (PacketError, ValueError):
        return {
            "status": STATUS_PACKET_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    try:
        restored_message = (
            parsed_packet.payload.decode(
                "utf-8"
            )
        )

    except UnicodeDecodeError:
        return {
            "status": STATUS_PACKET_ERROR,
            "marker_errors": marker_errors,
            "message": None,
        }

    if (
        parsed_packet.message_type
        != MESSAGE_TYPE_TEXT
        or parsed_packet.sequence_number
        != EXPECTED_SEQUENCE_NUMBER
        or restored_message
        != EXPECTED_MESSAGE
    ):
        return {
            "status": STATUS_PACKET_ERROR,
            "marker_errors": marker_errors,
            "message": restored_message,
        }

    return {
        "status": STATUS_SUCCESS,
        "marker_errors": marker_errors,
        "message": restored_message,
    }


# ============================================================
# Загрузка IQ
# ============================================================

if not INPUT_IQ_PATH.exists():
    raise FileNotFoundError(
        f"Не найден файл: {INPUT_IQ_PATH}. "
        "Сначала необходимо выполнить Lab018."
    )

transmitted_iq = np.load(
    INPUT_IQ_PATH
).astype(
    np.complex128
)

if transmitted_iq.ndim != 1:
    raise ValueError(
        "IQ-массив должен быть одномерным"
    )


# ============================================================
# Подготовка приёмника
# ============================================================

rrc_taps = root_raised_cosine_taps(
    rolloff=RRC_ROLLOFF,
    samples_per_symbol=SAMPLES_PER_SYMBOL,
    span_symbols=RRC_SPAN_SYMBOLS,
)

marker_bits, marker_symbols = (
    build_frame_marker()
)


# ============================================================
# Эксперимент
# ============================================================

results = []

worst_offset_diagnostic = None

for frequency_offset_hz in FREQUENCY_OFFSETS_HZ:

    impaired_iq = apply_frequency_offset(
        iq_samples=transmitted_iq,
        frequency_offset_hz=frequency_offset_hz,
        sample_rate=SAMPLE_RATE,
        phase_offset_degrees=(
            TEST_PHASE_OFFSET_DEGREES
        ),
    )

    received_iq = np.concatenate(
        [
            np.zeros(
                TEST_SAMPLE_DELAY,
                dtype=np.complex128,
            ),
            impaired_iq,
        ]
    )

    matched_iq = fftconvolve(
        received_iq,
        rrc_taps,
        mode="full",
    )

    search_result = (
        find_frame_and_frequency_offset(
            matched_iq=matched_iq,
            marker_symbols=marker_symbols,
            samples_per_symbol=(
                SAMPLES_PER_SYMBOL
            ),
            symbol_rate=SYMBOL_RATE,
        )
    )

    symbol_samples = search_result[
        "symbol_samples"
    ]

    frame_start_symbol = search_result[
        "start_symbol_index"
    ]

    estimated_frequency_hz = search_result[
        "frequency_offset_hz"
    ]

    estimated_initial_phase = search_result[
        "initial_phase"
    ]

    estimated_phase_increment = search_result[
        "phase_increment"
    ]

    symbol_indexes = np.arange(
        len(symbol_samples),
        dtype=np.float64,
    )

    relative_symbol_indexes = (
        symbol_indexes
        - frame_start_symbol
    )

    # --------------------------------------------------------
    # Приём без компенсации частотной ошибки
    # --------------------------------------------------------

    constant_phase_corrected = (
        symbol_samples
        * np.exp(
            -1j
            * estimated_initial_phase
        )
    )

    result_without_cfo = decode_radio_frame(
        corrected_symbols=(
            constant_phase_corrected
        ),
        frame_start_symbol=(
            frame_start_symbol
        ),
        marker_bits=marker_bits,
    )

    # --------------------------------------------------------
    # Полная фазовая и частотная коррекция
    # --------------------------------------------------------

    phase_model = (
        estimated_initial_phase
        + estimated_phase_increment
        * relative_symbol_indexes
    )

    frequency_corrected = (
        symbol_samples
        * np.exp(
            -1j * phase_model
        )
    )

    result_with_cfo = decode_radio_frame(
        corrected_symbols=(
            frequency_corrected
        ),
        frame_start_symbol=(
            frame_start_symbol
        ),
        marker_bits=marker_bits,
    )

    frequency_error_hz = (
        estimated_frequency_hz
        - frequency_offset_hz
    )

    results.append(
        {
            "true_frequency_hz": (
                frequency_offset_hz
            ),
            "estimated_frequency_hz": (
                estimated_frequency_hz
            ),
            "frequency_error_hz": (
                frequency_error_hz
            ),
            "correlation_score": (
                search_result["score"]
            ),
            "sample_phase": (
                search_result["sample_phase"]
            ),
            "frame_start_symbol": (
                frame_start_symbol
            ),
            "marker_errors_without": (
                result_without_cfo[
                    "marker_errors"
                ]
            ),
            "status_without": (
                result_without_cfo["status"]
            ),
            "marker_errors_with": (
                result_with_cfo[
                    "marker_errors"
                ]
            ),
            "status_with": (
                result_with_cfo["status"]
            ),
        }
    )

    if (
        frequency_offset_hz
        == max(FREQUENCY_OFFSETS_HZ)
    ):
        marker_start = frame_start_symbol
        marker_end = (
            marker_start
            + len(marker_symbols)
        )

        marker_before = (
            symbol_samples[
                marker_start:
                marker_end
            ]
            * marker_symbols
        )

        marker_after = (
            frequency_corrected[
                marker_start:
                marker_end
            ]
            * marker_symbols
        )

        worst_offset_diagnostic = {
            "before": marker_before,
            "after": marker_after,
        }


# ============================================================
# Вывод
# ============================================================

print(
    "=== Lab021. Компенсация частотной ошибки ==="
)

print("\nПараметры:")

print(
    "Символьная скорость:",
    SYMBOL_RATE,
    "символов/с",
)

print(
    "Частота дискретизации:",
    SAMPLE_RATE,
    "сэмплов/с",
)

print(
    "Постоянный фазовый поворот:",
    TEST_PHASE_OFFSET_DEGREES,
    "градусов",
)

print(
    "Задержка:",
    TEST_SAMPLE_DELAY,
    "сэмплов",
)


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

print(
    f"{'Δf задано':>12}"
    f"{'Δf оценено':>14}"
    f"{'Ошибка':>12}"
    f"{'Коррел.':>11}"
    f"{'Без CFO':>16}"
    f"{'С CFO':>16}"
)

print("-" * 81)

for result in results:

    print(
        f"{result['true_frequency_hz']:>9.1f} Гц"
        f"{result['estimated_frequency_hz']:>11.2f} Гц"
        f"{result['frequency_error_hz']:>9.2f} Гц"
        f"{result['correlation_score']:>11.4f}"
        f"{result['status_without']:>16}"
        f"{result['status_with']:>16}"
    )


# ============================================================
# Графики
# ============================================================

true_frequencies = np.array(
    [
        result["true_frequency_hz"]
        for result in results
    ],
    dtype=np.float64,
)

estimated_frequencies = np.array(
    [
        result["estimated_frequency_hz"]
        for result in results
    ],
    dtype=np.float64,
)

marker_errors_without = np.array(
    [
        (
            result["marker_errors_without"]
            if result["marker_errors_without"]
            is not None
            else len(marker_bits)
        )
        for result in results
    ],
    dtype=np.float64,
)

marker_errors_with = np.array(
    [
        (
            result["marker_errors_with"]
            if result["marker_errors_with"]
            is not None
            else len(marker_bits)
        )
        for result in results
    ],
    dtype=np.float64,
)

success_without = np.array(
    [
        result["status_without"]
        == STATUS_SUCCESS
        for result in results
    ],
    dtype=np.int32,
)

success_with = np.array(
    [
        result["status_with"]
        == STATUS_SUCCESS
        for result in results
    ],
    dtype=np.int32,
)


figure, axes = plt.subplots(
    2,
    2,
    figsize=(13, 10),
)


# ------------------------------------------------------------
# 1. Оценка частоты
# ------------------------------------------------------------

axes[0, 0].plot(
    true_frequencies,
    true_frequencies,
    linestyle="--",
    label="Идеальная оценка",
)

axes[0, 0].plot(
    true_frequencies,
    estimated_frequencies,
    marker="o",
    label="Оценка приёмника",
)

axes[0, 0].set_xlabel(
    "Заданное смещение, Гц"
)

axes[0, 0].set_ylabel(
    "Оценённое смещение, Гц"
)

axes[0, 0].set_title(
    "Оценка частотного рассогласования"
)

axes[0, 0].grid(
    True
)

axes[0, 0].legend()


# ------------------------------------------------------------
# 2. Ошибки маркера
# ------------------------------------------------------------

axes[0, 1].plot(
    true_frequencies,
    marker_errors_without,
    marker="o",
    label="Без компенсации CFO",
)

axes[0, 1].plot(
    true_frequencies,
    marker_errors_with,
    marker="s",
    label="После компенсации CFO",
)

axes[0, 1].set_xlabel(
    "Частотное смещение, Гц"
)

axes[0, 1].set_ylabel(
    "Ошибки PREAMBLE + SYNC"
)

axes[0, 1].set_title(
    "Ошибки известного маркера"
)

axes[0, 1].grid(
    True
)

axes[0, 1].legend()


# ------------------------------------------------------------
# 3. Успешный приём кадра
# ------------------------------------------------------------

axes[1, 0].plot(
    true_frequencies,
    success_without,
    marker="o",
    label="Без компенсации CFO",
)

axes[1, 0].plot(
    true_frequencies,
    success_with,
    marker="s",
    label="После компенсации CFO",
)

axes[1, 0].set_yticks(
    [0, 1],
    [
        "Ошибка",
        "SUCCESS",
    ],
)

axes[1, 0].set_xlabel(
    "Частотное смещение, Гц"
)

axes[1, 0].set_ylabel(
    "Результат приёма"
)

axes[1, 0].set_title(
    "Восстановление полного пакета"
)

axes[1, 0].grid(
    True
)

axes[1, 0].legend()


# ------------------------------------------------------------
# 4. Фаза маркера при максимальном CFO
# ------------------------------------------------------------

if worst_offset_diagnostic is None:
    raise RuntimeError(
        "Не сохранена диагностика максимального CFO"
    )

phase_before = np.unwrap(
    np.angle(
        worst_offset_diagnostic["before"]
    )
)

phase_after = np.unwrap(
    np.angle(
        worst_offset_diagnostic["after"]
    )
)

axes[1, 1].plot(
    phase_before,
    label="До компенсации",
)

axes[1, 1].plot(
    phase_after,
    label="После компенсации",
)

axes[1, 1].set_xlabel(
    "Номер символа маркера"
)

axes[1, 1].set_ylabel(
    "Развёрнутая фаза, рад"
)

axes[1, 1].set_title(
    "Фаза при Δf = "
    f"{max(FREQUENCY_OFFSETS_HZ):.0f} Гц"
)

axes[1, 1].grid(
    True
)

axes[1, 1].legend()


figure.tight_layout()

figure.savefig(
    GRAPH_PATH,
    dpi=160,
)

plt.close(
    figure
)


# ============================================================
# Отчёт
# ============================================================

report_lines = [
    "Lab021. Frequency offset correction",
    "",
    f"Symbol rate: {SYMBOL_RATE}",
    f"Sample rate: {SAMPLE_RATE}",
    (
        "Phase offset: "
        f"{TEST_PHASE_OFFSET_DEGREES:.2f} deg"
    ),
    (
        "Sample delay: "
        f"{TEST_SAMPLE_DELAY}"
    ),
    "",
]

for result in results:

    report_lines.extend(
        [
            (
                "True frequency offset: "
                f"{result['true_frequency_hz']:.2f} Hz"
            ),
            (
                "Estimated frequency offset: "
                f"{result['estimated_frequency_hz']:.4f} Hz"
            ),
            (
                "Frequency estimation error: "
                f"{result['frequency_error_hz']:.4f} Hz"
            ),
            (
                "Correlation score: "
                f"{result['correlation_score']:.6f}"
            ),
            (
                "Status without correction: "
                f"{result['status_without']}"
            ),
            (
                "Status with correction: "
                f"{result['status_with']}"
            ),
            "",
        ]
    )

REPORT_PATH.write_text(
    "\n".join(
        report_lines
    ),
    encoding="utf-8",
)


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

result_at_zero = next(
    result
    for result in results
    if result["true_frequency_hz"] == 0.0
)

result_at_maximum = next(
    result
    for result in results
    if (
        result["true_frequency_hz"]
        == max(FREQUENCY_OFFSETS_HZ)
    )
)

assert (
    result_at_zero["status_without"]
    == STATUS_SUCCESS
)

assert all(
    result["status_with"] == STATUS_SUCCESS
    for result in results
)

assert abs(
    result_at_maximum["frequency_error_hz"]
) < 5.0

assert (
    result_at_maximum["status_without"]
    != STATUS_SUCCESS
)

assert GRAPH_PATH.exists()
assert REPORT_PATH.exists()


print("\nГрафик:")

print(GRAPH_PATH)

print("\nОтчёт:")

print(REPORT_PATH)

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