"""
Run a custom Ultralytics YOLO26 segmentation model
on shadowgraph flow images.

For each input frame, the script saves:
    1. An image with overlaid masks in the RESULTS_DIRECTORY directory.
    2. The same image with the fitted median shock-wave-front line
       in the RESULTS_WITH_LINES_DIRECTORY directory.
    3. Fitted lines on a coordinate grid, with their equations shown
       in the legend, in the RESULTS_LINES_ON_GRID_DIRECTORY directory.

Install the dependencies in the PyCharm terminal:
    python -m pip install -U ultralytics opencv-python numpy matplotlib
"""

from pathlib import Path
import importlib
import os
import re
import shlex
import subprocess
import sys
import time
import warnings

try:
    import cv2
    import numpy as np
except ImportError as exc:
    raise SystemExit(
        "Не установлены OpenCV и/или NumPy. Выполните команду:\n"
        "python -m pip install -U opencv-python numpy"
    ) from exc

try:
    import matplotlib

    # No interactive plot window is needed because all figures are saved
    # to files. The Agg backend also works in headless environments.
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.ticker import MultipleLocator
except ImportError as exc:
    raise SystemExit(
        "Не установлен Matplotlib. Выполните команду:\n"
        "python -m pip install -U matplotlib"
    ) from exc


# =============================================================================
# USER PARAMETERS — change settings only in this section
# =============================================================================

# Paths can be absolute or specified relative to this script.
modelName = "shock_seg_model_YOLO26.pt"
framesDirectory = "frames"
resultsDirectory = "results"
resultsWithLinesDirectory = "results_with_lines"
resultsLinesOnGridDirectory = "results_lines_on_grid"

# YOLO26 compatibility. Ultralytics 8.4.0 is the minimum version for YOLO26
# family models. If the package is missing, outdated, or does not contain the
# C3k2 layer, the script upgrades it in the current virtual environment and restarts.
autoUpdateUltralytics = True
minimumUltralyticsVersion = "8.4.0"

# Time displayed on the images. The first frame after alphabetical sorting
# has index 0:
# frame_time = firstFrameTime + frame_index * timeStep
printTime = True
firstFrameTime = 0.0
timeStep = 6.67
timeUnit = "us"
timeDecimalPlaces = 2

# YOLO inference parameters.
confidenceThreshold = 0.25
iouThreshold = 0.70
imageSize = 640
device = None                 # None selects automatically; examples: "cpu", 0, "mps"
useHalfPrecision = False      # True is valid only for a compatible CUDA GPU.

# Mask visualization parameters.
maskThreshold = 0.50
maskOpacity = 0.35
maskContourThickness = 2
drawMaskLabels = True
minimumMaskArea = 25          # Masks with a smaller area (in pixels) are ignored.
keepLargestConnectedComponent = True

# Median shock-wave-front line.
# "auto": y=f(x) for a predominantly horizontal mask and x=f(y)
# for a predominantly vertical mask.
# Other valid values are "y_of_x" and "x_of_y".
centerlineAxis = "auto"
polynomialDegree = 2          # 1 = line, 2 = parabola, 3 = cubic curve, and so on.
centerlineSamplingStep = 1
minimumCenterlinePoints = 8
robustFitIterations = 3
outlierSigma = 3.5
centerlineThickness = 1

# Separate median-line plots on a coordinate grid.
# Coordinates and equations are expressed in pixels of the original image.
gridFigureWidthInches = 10.0
gridFigureMinimumHeightInches = 4.0
gridFigureMaximumHeightInches = 12.0
gridFigureDpi = 150
gridMajorStep = 50
gridMinorStep = 10
gridLineWidth = 2.0
gridEquationSignificantDigits = 4
gridLegendFontSize = 8
gridAxisFontSize = 9
gridTitleFontSize = 9
gridInvertYAxis = True       # True: the coordinate origin matches the frame's top-left corner.
gridBackgroundColor = "white"
gridMajorColor = "#9a9a9a"
gridMinorColor = "#d8d8d8"

# Text and output-image parameters.
fontScale = 0.38
fontThickness = 1
textMargin = 6
textLineSpacing = 4
showNoMasksMessage = True
jpegQuality = 95
pngCompression = 3

# =============================================================================
# END OF USER PARAMETERS
# =============================================================================


SCRIPT_DIRECTORY = Path(__file__).resolve().parent
SUPPORTED_EXTENSIONS = {".png", ".jpg", ".jpeg"}
ULTRALYTICS_RESTART_FLAG = "SHOCK_SEG_ULTRALYTICS_RESTARTED"

# BGR colors cycled when multiple shock-wave masks are detected.
MASK_COLORS = [
    (0, 0, 255),       # red
    (0, 165, 255),     # orange
    (0, 255, 255),     # yellow
    (255, 0, 255),     # magenta
    (255, 255, 0),     # cyan
    (0, 255, 0),       # green
    (255, 128, 0),     # blue-orange
]


def version_tuple(version_text):
    """Convert a version number such as 8.4.115 into a comparable integer tuple."""
    numbers = [int(value) for value in re.findall(r"\d+", str(version_text))[:3]]
    return tuple(numbers + [0] * (3 - len(numbers)))


def pip_update_command(force_reinstall=False):
    """Build the update command for the current Python interpreter."""
    command = [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--upgrade",
    ]
    if force_reinstall:
        # Reinstall only the Ultralytics package; already installed large
        # dependencies such as PyTorch are not downloaded again.
        command.extend(["--force-reinstall", "--no-deps"])
    command.append(f"ultralytics>={minimumUltralyticsVersion}")
    return command


def printable_command(command):
    """Return a command string that is convenient to copy into a terminal."""
    return shlex.join(command)


def inspect_ultralytics():
    """Check the Ultralytics version and the availability of the C3k2 class."""
    problems = []
    force_reinstall = False

    try:
        ultralytics_module = importlib.import_module("ultralytics")
    except ModuleNotFoundError as exc:
        if exc.name == "ultralytics":
            return None, None, ["пакет ultralytics не установлен"], False
        return None, None, [f"ошибка импорта зависимости: {exc}"], False
    except Exception as exc:
        return None, None, [f"ошибка импорта ultralytics: {exc}"], True

    installed_version = getattr(ultralytics_module, "__version__", "неизвестно")
    version_is_old = (
        version_tuple(installed_version)
        < version_tuple(minimumUltralyticsVersion)
    )
    if version_is_old:
        problems.append(
            f"установлена версия {installed_version}, требуется не ниже "
            f"{minimumUltralyticsVersion}"
        )

    try:
        block_module = importlib.import_module("ultralytics.nn.modules.block")
        if not hasattr(block_module, "C3k2"):
            problems.append("в пакете отсутствует слой C3k2")
            # For an old version, a regular upgrade with dependencies is enough.
            # A forced reinstall is needed only when the version number is
            # already sufficient but the package files are corrupted.
            force_reinstall = not version_is_old
    except Exception as exc:
        problems.append(f"не удалось импортировать слой C3k2: {exc}")
        force_reinstall = not version_is_old

    return ultralytics_module, installed_version, problems, force_reinstall


def get_yolo_class():
    """
    Prepare a compatible Ultralytics version and return the YOLO class.

    The upgrade is performed through sys.executable, so it runs inside the
    environment selected in the PyCharm run configuration.
    """
    module, installed_version, problems, force_reinstall = inspect_ultralytics()

    if problems:
        details = "; ".join(problems)
        command = pip_update_command(force_reinstall=force_reinstall)
        command_text = printable_command(command)

        if not autoUpdateUltralytics:
            raise RuntimeError(
                "Текущая установка Ultralytics несовместима с моделью YOLO26: "
                f"{details}.\nВыполните команду:\n{command_text}"
            )

        if os.environ.get(ULTRALYTICS_RESTART_FLAG) == "1":
            raise RuntimeError(
                "После автоматического обновления Ultralytics проблема "
                f"сохранилась: {details}.\nВыполните вручную:\n{command_text}\n"
                "Если модель обучалась в изменённой версии Ultralytics, "
                "установите ту же версию, которая использовалась при обучении."
            )

        print(
            "Обнаружена несовместимая установка Ultralytics: "
            f"{details}.\nВыполняется обновление:\n{command_text}"
        )
        completed = subprocess.run(command, check=False)
        if completed.returncode != 0:
            raise RuntimeError(
                "Не удалось автоматически обновить Ultralytics. "
                f"Выполните команду вручную:\n{command_text}"
            )

        # Already imported modules cannot be replaced reliably within the same
        # process, so the script performs one clean restart after the upgrade.
        os.environ[ULTRALYTICS_RESTART_FLAG] = "1"
        script_path = str(Path(__file__).resolve())
        os.execv(sys.executable, [sys.executable, script_path, *sys.argv[1:]])

    os.environ.pop(ULTRALYTICS_RESTART_FLAG, None)
    try:
        return module.YOLO
    except AttributeError as exc:
        raise RuntimeError(
            f"В Ultralytics {installed_version} отсутствует класс YOLO. "
            "Переустановите пакет Ultralytics."
        ) from exc


def load_yolo_model(yolo_class, model_path):
    """Load the weights and provide a precise explanation for deserialization errors."""
    try:
        return yolo_class(str(model_path))
    except AttributeError as exc:
        if "Can't get attribute" in str(exc):
            raise RuntimeError(
                "Не удалось десериализовать модель. Установленная версия "
                "Ultralytics не содержит архитектурный класс из файла весов. "
                f"Исходная ошибка: {exc}\n"
                "Обновите Ultralytics либо установите точную версию пакета, "
                "в которой модель была обучена."
            ) from exc
        raise


def resolve_path(path_value):
    """Resolve a user-provided path relative to the script directory."""
    path = Path(path_value).expanduser()
    return path if path.is_absolute() else SCRIPT_DIRECTORY / path


def read_image(path):
    """Read an image, including from a path containing non-ASCII characters."""
    try:
        encoded = np.fromfile(str(path), dtype=np.uint8)
        image = cv2.imdecode(encoded, cv2.IMREAD_COLOR)
    except (OSError, ValueError):
        image = None

    if image is None:
        raise RuntimeError(f"Could not read image: {path}")
    return image


def write_image(path, image):
    """Save an image, including to a path containing non-ASCII characters."""
    suffix = path.suffix.lower()
    if suffix in {".jpg", ".jpeg"}:
        parameters = [cv2.IMWRITE_JPEG_QUALITY, int(jpegQuality)]
    elif suffix == ".png":
        parameters = [cv2.IMWRITE_PNG_COMPRESSION, int(pngCompression)]
    else:
        raise ValueError(f"Unsupported output extension: {suffix}")

    success, encoded = cv2.imencode(suffix, image, parameters)
    if not success:
        raise RuntimeError(f"OpenCV could not encode image: {path}")

    try:
        encoded.tofile(str(path))
    except OSError as exc:
        raise RuntimeError(f"Could not save image: {path}") from exc


def find_frames(directory):
    """Return PNG/JPG/JPEG files in case-insensitive alphabetical order."""
    frames = [
        path
        for path in directory.iterdir()
        if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS
    ]
    return sorted(frames, key=lambda path: (path.name.casefold(), path.name))


def largest_connected_component(binary_mask):
    """Keep only the largest connected component of a binary mask."""
    mask_uint8 = binary_mask.astype(np.uint8)
    number, labels, statistics, _ = cv2.connectedComponentsWithStats(
        mask_uint8, connectivity=8
    )

    if number <= 1:
        return binary_mask

    component_areas = statistics[1:, cv2.CC_STAT_AREA]
    largest_label = 1 + int(np.argmax(component_areas))
    return labels == largest_label


def get_class_name(names, class_id):
    """Get a readable class name from the Ultralytics Results.names object."""
    if isinstance(names, dict):
        return str(names.get(class_id, class_id))
    if isinstance(names, (list, tuple)) and 0 <= class_id < len(names):
        return str(names[class_id])
    return str(class_id)


def extract_detections(result, image_shape):
    """Convert predicted Ultralytics masks into CPU NumPy arrays."""
    if result.masks is None or result.masks.data is None:
        return []

    masks_data = result.masks.data
    if hasattr(masks_data, "cpu"):
        masks_data = masks_data.cpu().numpy()
    else:
        masks_data = np.asarray(masks_data)

    confidences = np.ones(len(masks_data), dtype=float)
    class_ids = np.zeros(len(masks_data), dtype=int)

    if result.boxes is not None:
        if result.boxes.conf is not None:
            confidences = result.boxes.conf.detach().cpu().numpy().astype(float)
        if result.boxes.cls is not None:
            class_ids = result.boxes.cls.detach().cpu().numpy().astype(int)

    image_height, image_width = image_shape[:2]
    detections = []

    for index, predicted_mask in enumerate(masks_data):
        if predicted_mask.shape != (image_height, image_width):
            predicted_mask = cv2.resize(
                predicted_mask.astype(np.float32),
                (image_width, image_height),
                interpolation=cv2.INTER_LINEAR,
            )

        binary_mask = predicted_mask >= maskThreshold
        if keepLargestConnectedComponent and np.any(binary_mask):
            binary_mask = largest_connected_component(binary_mask)

        area = int(np.count_nonzero(binary_mask))
        if area < minimumMaskArea:
            continue

        class_id = int(class_ids[index]) if index < len(class_ids) else 0
        confidence = float(confidences[index]) if index < len(confidences) else 1.0
        detections.append(
            {
                "mask": binary_mask,
                "class_id": class_id,
                "class_name": get_class_name(result.names, class_id),
                "confidence": confidence,
                "area": area,
            }
        )

    return detections


def draw_label(image, text, x, y, color):
    """Draw one clearly readable label on a dark background."""
    (text_width, text_height), baseline = cv2.getTextSize(
        text, cv2.FONT_HERSHEY_SIMPLEX, fontScale, fontThickness
    )
    x = int(np.clip(x, 0, max(0, image.shape[1] - text_width - 7)))
    y = int(np.clip(y, text_height + 6, max(text_height + 6, image.shape[0] - 2)))

    cv2.rectangle(
        image,
        (x, y - text_height - 6),
        (x + text_width + 7, y + baseline + 3),
        (20, 20, 20),
        thickness=-1,
    )
    cv2.putText(
        image,
        text,
        (x + 3, y - 2),
        cv2.FONT_HERSHEY_SIMPLEX,
        fontScale,
        color,
        fontThickness,
        cv2.LINE_AA,
    )


def draw_masks(image, detections):
    """Overlay all detected masks, contours, and optional class labels."""
    output = image.copy()

    for index, detection in enumerate(detections):
        mask = detection["mask"]
        color = MASK_COLORS[index % len(MASK_COLORS)]

        color_layer = np.empty_like(output)
        color_layer[:] = color
        blended = cv2.addWeighted(output, 1.0 - maskOpacity, color_layer, maskOpacity, 0)
        output[mask] = blended[mask]

        contours, _ = cv2.findContours(
            mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        cv2.drawContours(output, contours, -1, color, maskContourThickness, cv2.LINE_AA)

        if drawMaskLabels:
            y_pixels, x_pixels = np.nonzero(mask)
            x_min = int(x_pixels.min())
            y_min = int(y_pixels.min())
            label = (
                f"{detection['class_name']} {detection['confidence']:.2f}"
            )
            draw_label(output, label, x_min, max(y_min - 4, 0), color)

    return output


def choose_centerline_axis(mask):
    """Choose the y=f(x) or x=f(y) representation of the front."""
    if centerlineAxis not in {"auto", "y_of_x", "x_of_y"}:
        raise ValueError(
            'centerlineAxis must be "auto", "y_of_x", or "x_of_y".'
        )
    if centerlineAxis != "auto":
        return centerlineAxis

    y_pixels, x_pixels = np.nonzero(mask)
    x_span = int(np.ptp(x_pixels))
    y_span = int(np.ptp(y_pixels))
    return "y_of_x" if x_span >= y_span else "x_of_y"


def median_centerline_samples(mask, axis):
    """Calculate median coordinates across the thickness of one mask."""
    y_pixels, x_pixels = np.nonzero(mask)

    if axis == "y_of_x":
        independent_all = x_pixels
        dependent_all = y_pixels
    else:
        independent_all = y_pixels
        dependent_all = x_pixels

    independent_values = np.unique(independent_all)
    step = max(1, int(centerlineSamplingStep))
    independent_values = independent_values[::step]

    dependent_medians = np.array(
        [
            np.median(dependent_all[independent_all == coordinate])
            for coordinate in independent_values
        ],
        dtype=np.float64,
    )
    return independent_values.astype(np.float64), dependent_medians


def robust_polynomial_fit(independent, dependent, requested_degree):
    """Fit a polynomial with iterative outlier rejection."""
    if len(independent) < 2:
        return None

    degree = min(max(1, int(requested_degree)), len(independent) - 1)
    valid = np.isfinite(independent) & np.isfinite(dependent)

    if np.count_nonzero(valid) < max(minimumCenterlinePoints, degree + 1):
        return None

    coefficients = None
    for _ in range(max(1, int(robustFitIterations))):
        if np.count_nonzero(valid) < degree + 1:
            break

        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            coefficients = np.polyfit(
                independent[valid], dependent[valid], degree
            )

        residuals = dependent - np.polyval(coefficients, independent)
        residual_center = np.median(residuals[valid])
        mad = np.median(np.abs(residuals[valid] - residual_center))
        robust_standard_deviation = 1.4826 * mad

        if robust_standard_deviation <= 1e-12:
            break

        updated_valid = (
            np.abs(residuals - residual_center)
            <= float(outlierSigma) * robust_standard_deviation
        )
        if np.count_nonzero(updated_valid) < degree + 1 or np.array_equal(
            updated_valid, valid
        ):
            break
        valid = updated_valid

    return coefficients


def fit_centerline(mask):
    """Fit the median line and return its points, axis, and coefficients."""
    axis = choose_centerline_axis(mask)
    independent, dependent = median_centerline_samples(mask, axis)
    coefficients = robust_polynomial_fit(
        independent, dependent, polynomialDegree
    )
    if coefficients is None:
        return None

    independent_line = np.arange(
        int(np.ceil(independent.min())),
        int(np.floor(independent.max())) + 1,
        dtype=np.float64,
    )
    dependent_line = np.polyval(coefficients, independent_line)

    if axis == "y_of_x":
        x_line = independent_line
        y_line = dependent_line
    else:
        x_line = dependent_line
        y_line = independent_line

    image_height, image_width = mask.shape
    valid = (
        np.isfinite(x_line)
        & np.isfinite(y_line)
        & (x_line >= 0)
        & (x_line < image_width)
        & (y_line >= 0)
        & (y_line < image_height)
    )
    plot_points = np.column_stack((x_line[valid], y_line[valid]))
    points = np.rint(plot_points).astype(np.int32)

    if len(points) < 2:
        return None
    return {
        "points": points,
        "plot_points": plot_points,
        "axis": axis,
        "coefficients": coefficients,
    }


def format_number_for_equation(value):
    """Format a coefficient compactly for mathematical notation."""
    value = float(abs(value))
    if value == 0.0:
        return "0"

    digits = max(1, int(gridEquationSignificantDigits))
    exponent = int(np.floor(np.log10(value)))

    # Scientific notation remains readable for the very small and large
    # coefficients that often occur when fitting in pixel coordinates.
    if exponent <= -3 or exponent >= 4:
        mantissa = value / (10.0 ** exponent)
        mantissa_text = f"{mantissa:.{digits - 1}f}".rstrip("0").rstrip(".")
        return rf"{mantissa_text}\times 10^{{{exponent}}}"

    decimal_places = max(0, digits - exponent - 1)
    return f"{value:.{decimal_places}f}".rstrip("0").rstrip(".")


def format_polynomial_equation(axis, coefficients):
    """Build a polynomial equation for display in the Matplotlib legend."""
    dependent_variable = "y" if axis == "y_of_x" else "x"
    independent_variable = "x" if axis == "y_of_x" else "y"
    degree = len(coefficients) - 1
    terms = []

    for coefficient_index, coefficient in enumerate(coefficients):
        power = degree - coefficient_index
        coefficient = float(coefficient)
        if abs(coefficient) < 1e-15:
            continue

        magnitude_text = format_number_for_equation(coefficient)
        if power == 0:
            variable_text = ""
        elif power == 1:
            variable_text = independent_variable
        else:
            variable_text = rf"{independent_variable}^{{{power}}}"

        term_text = magnitude_text + variable_text
        if not terms:
            terms.append(("-" if coefficient < 0 else "") + term_text)
        else:
            terms.append((" - " if coefficient < 0 else " + ") + term_text)

    expression = "".join(terms) if terms else "0"
    return rf"${dependent_variable} = {expression}$"


def save_centerlines_on_grid(output_path, centerline_fits, image_shape, frame_index):
    """Save lines on a coordinate grid with their equations in the legend."""
    image_height, image_width = image_shape[:2]
    aspect_ratio = image_height / max(1, image_width)
    figure_height = float(
        np.clip(
            float(gridFigureWidthInches) * aspect_ratio,
            float(gridFigureMinimumHeightInches),
            float(gridFigureMaximumHeightInches),
        )
    )

    figure, axes = plt.subplots(
        figsize=(float(gridFigureWidthInches), figure_height),
        dpi=int(gridFigureDpi),
    )
    axes.set_facecolor(gridBackgroundColor)

    for centerline_index, (centerline_fit, bgr_color) in enumerate(
        centerline_fits, start=1
    ):
        points = centerline_fit["plot_points"]
        rgb_color = tuple(channel / 255.0 for channel in reversed(bgr_color))
        equation = format_polynomial_equation(
            centerline_fit["axis"], centerline_fit["coefficients"]
        )
        axes.plot(
            points[:, 0],
            points[:, 1],
            color=rgb_color,
            linewidth=float(gridLineWidth),
            label=f"Фронт {centerline_index}: {equation}",
        )

    axes.set_xlim(0, max(1, image_width - 1))
    if gridInvertYAxis:
        axes.set_ylim(max(1, image_height - 1), 0)
    else:
        axes.set_ylim(0, max(1, image_height - 1))
    axes.set_aspect("equal", adjustable="box")
    axes.set_xlabel("x, пикс.", fontsize=gridAxisFontSize)
    axes.set_ylabel("y, пикс.", fontsize=gridAxisFontSize)
    axes.tick_params(axis="both", which="major", labelsize=gridAxisFontSize)

    axes.xaxis.set_major_locator(MultipleLocator(float(gridMajorStep)))
    axes.yaxis.set_major_locator(MultipleLocator(float(gridMajorStep)))
    axes.xaxis.set_minor_locator(MultipleLocator(float(gridMinorStep)))
    axes.yaxis.set_minor_locator(MultipleLocator(float(gridMinorStep)))
    axes.grid(which="major", color=gridMajorColor, linewidth=0.8, alpha=0.85)
    axes.grid(which="minor", color=gridMinorColor, linewidth=0.45, alpha=0.8)

    if printTime:
        current_time = firstFrameTime + frame_index * timeStep
        axes.set_title(
            f"t = {current_time:.{timeDecimalPlaces}f} {timeUnit}",
            loc="left",
            fontsize=gridTitleFontSize,
        )

    if centerline_fits:
        axes.legend(
            loc="best",
            fontsize=gridLegendFontSize,
            framealpha=0.92,
            facecolor="white",
            edgecolor="#555555",
        )
    elif showNoMasksMessage:
        axes.text(
            0.5,
            0.5,
            "Линии ударных волн не обнаружены",
            transform=axes.transAxes,
            horizontalalignment="center",
            verticalalignment="center",
            fontsize=gridAxisFontSize,
            color="#333333",
        )

    figure.tight_layout()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    save_arguments = {"dpi": int(gridFigureDpi), "facecolor": "white"}
    if output_path.suffix.lower() in {".jpg", ".jpeg"}:
        save_arguments["pil_kwargs"] = {"quality": int(jpegQuality)}
    figure.savefig(output_path, **save_arguments)
    plt.close(figure)


def fit_text_scale(lines, image_width):
    """Reduce the overall label scale when a line extends beyond the frame."""
    scale = float(fontScale)
    available_width = max(1, image_width - 2 * int(textMargin) - 12)

    for text, _ in lines:
        (width, _), _ = cv2.getTextSize(
            text, cv2.FONT_HERSHEY_SIMPLEX, scale, fontThickness
        )
        if width > available_width:
            scale *= available_width / width

    return max(0.28, scale)


def draw_text_block(image, lines):
    """Draw colored lines in a translucent block in the top-left corner."""
    if not lines:
        return

    scale = fit_text_scale(lines, image.shape[1])
    sizes = [
        cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, scale, fontThickness)
        for text, _ in lines
    ]
    line_heights = [size[0][1] + size[1] for size in sizes]
    block_width = max(size[0][0] for size in sizes) + 2 * textMargin
    block_height = (
        sum(line_heights)
        + (len(lines) - 1) * textLineSpacing
        + 2 * textMargin
    )

    overlay = image.copy()
    cv2.rectangle(
        overlay,
        (0, 0),
        (min(block_width, image.shape[1] - 1), min(block_height, image.shape[0] - 1)),
        (15, 15, 15),
        thickness=-1,
    )
    cv2.addWeighted(overlay, 0.72, image, 0.28, 0, image)

    y = textMargin
    for (text, color), ((_, text_height), baseline) in zip(lines, sizes):
        y += text_height
        cv2.putText(
            image,
            text,
            (textMargin, y),
            cv2.FONT_HERSHEY_SIMPLEX,
            scale,
            color,
            fontThickness,
            cv2.LINE_AA,
        )
        y += baseline + textLineSpacing


def add_time_annotation(lines, frame_index):
    """Add the current frame time to the labels when time display is enabled."""
    if printTime:
        current_time = firstFrameTime + frame_index * timeStep
        time_text = f"t = {current_time:.{timeDecimalPlaces}f} {timeUnit}"
        lines.append((time_text, (255, 255, 255)))


def process_frame(model, frame_path, frame_index):
    """Run inference and save three visualization variants."""
    image = read_image(frame_path)

    predict_arguments = {
        "source": image,
        "conf": confidenceThreshold,
        "iou": iouThreshold,
        "imgsz": imageSize,
        "retina_masks": True,
        "verbose": False,
        "save": False,
    }
    if device is not None:
        predict_arguments["device"] = device
    if useHalfPrecision:
        predict_arguments["half"] = True

    predictions = model.predict(**predict_arguments)
    if not predictions:
        raise RuntimeError(f"The model returned no Results object for {frame_path.name}")

    result = predictions[0]
    detections = extract_detections(result, image.shape)
    mask_image = draw_masks(image, detections)
    lines_image = mask_image.copy()

    mask_text_lines = []
    add_time_annotation(mask_text_lines, frame_index)

    line_text_lines = []
    add_time_annotation(line_text_lines, frame_index)

    centerline_fits = []

    for detection_index, detection in enumerate(detections):
        color = MASK_COLORS[detection_index % len(MASK_COLORS)]
        centerline_fit = fit_centerline(detection["mask"])
        if centerline_fit is None:
            continue

        points = centerline_fit["points"]
        centerline_fits.append((centerline_fit, color))

        cv2.polylines(
            lines_image,
            [points.reshape(-1, 1, 2)],
            isClosed=False,
            color=color,
            thickness=centerlineThickness,
            lineType=cv2.LINE_AA,
        )

    if not detections and showNoMasksMessage:
        no_mask_line = ("No shock masks detected", (255, 255, 255))
        mask_text_lines.append(no_mask_line)
        line_text_lines.append(no_mask_line)

    draw_text_block(mask_image, mask_text_lines)
    draw_text_block(lines_image, line_text_lines)

    write_image(resolve_path(resultsDirectory) / frame_path.name, mask_image)
    write_image(
        resolve_path(resultsWithLinesDirectory) / frame_path.name,
        lines_image,
    )
    save_centerlines_on_grid(
        resolve_path(resultsLinesOnGridDirectory) / frame_path.name,
        centerline_fits,
        image.shape,
        frame_index,
    )
    return len(detections)


def validate_parameters():
    """Validate configuration parameter values before processing begins."""
    if not 0.0 <= maskOpacity <= 1.0:
        raise ValueError("maskOpacity must be between 0 and 1.")
    if not 0.0 <= maskThreshold <= 1.0:
        raise ValueError("maskThreshold must be between 0 and 1.")
    if not 0.0 <= confidenceThreshold <= 1.0:
        raise ValueError("confidenceThreshold must be between 0 and 1.")
    if imageSize <= 0:
        raise ValueError("imageSize must be positive.")
    if polynomialDegree < 1:
        raise ValueError("polynomialDegree must be at least 1.")
    if gridFigureWidthInches <= 0:
        raise ValueError("gridFigureWidthInches must be positive.")
    if gridFigureMinimumHeightInches <= 0 or gridFigureMaximumHeightInches <= 0:
        raise ValueError("Grid figure height limits must be positive.")
    if gridFigureMinimumHeightInches > gridFigureMaximumHeightInches:
        raise ValueError(
            "gridFigureMinimumHeightInches cannot exceed "
            "gridFigureMaximumHeightInches."
        )
    if gridFigureDpi <= 0:
        raise ValueError("gridFigureDpi must be positive.")
    if gridMajorStep <= 0 or gridMinorStep <= 0:
        raise ValueError("gridMajorStep and gridMinorStep must be positive.")
    if gridEquationSignificantDigits < 1:
        raise ValueError("gridEquationSignificantDigits must be at least 1.")


def main():
    validate_parameters()

    model_path = resolve_path(modelName)
    frames_path = resolve_path(framesDirectory)
    results_path = resolve_path(resultsDirectory)
    results_with_lines_path = resolve_path(resultsWithLinesDirectory)
    results_lines_on_grid_path = resolve_path(resultsLinesOnGridDirectory)

    if not model_path.is_file():
        raise FileNotFoundError(
            f"Model file was not found: {model_path}\n"
            "Place the .pt file next to this script or change modelName."
        )
    if not frames_path.is_dir():
        raise NotADirectoryError(
            f"Frames directory was not found: {frames_path}\n"
            "Create it next to this script or change framesDirectory."
        )

    frames = find_frames(frames_path)
    if not frames:
        raise FileNotFoundError(
            f"No PNG/JPG/JPEG files were found in: {frames_path}"
        )

    results_path.mkdir(parents=True, exist_ok=True)
    results_with_lines_path.mkdir(parents=True, exist_ok=True)
    results_lines_on_grid_path.mkdir(parents=True, exist_ok=True)

    yolo_class = get_yolo_class()

    print(f"Загрузка модели: {model_path}")
    model = load_yolo_model(yolo_class, model_path)
    print(f"Found {len(frames)} frame(s). Processing alphabetically...")

    started_at = time.perf_counter()
    processed_count = 0
    failed_count = 0

    for frame_index, frame_path in enumerate(frames):
        frame_started_at = time.perf_counter()
        try:
            detection_count = process_frame(
                model, frame_path, frame_index
            )
            processed_count += 1
            elapsed = time.perf_counter() - frame_started_at
            print(
                f"[{frame_index + 1:04d}/{len(frames):04d}] "
                f"{frame_path.name} -> {detection_count} mask(s), {elapsed:.3f} s"
            )
        except Exception as exc:  # A corrupted frame must not stop the batch.
            failed_count += 1
            print(
                f"[{frame_index + 1:04d}/{len(frames):04d}] "
                f"ERROR in {frame_path.name}: {exc}",
                file=sys.stderr,
            )

    total_elapsed = time.perf_counter() - started_at
    print("\nFinished.")
    print(f"Successfully processed: {processed_count}")
    print(f"Failed: {failed_count}")
    print(f"Total time: {total_elapsed:.2f} s")
    print(f"Mask results: {results_path}")
    print(f"Mask + line results: {results_with_lines_path}")
    print(f"Lines on grid: {results_lines_on_grid_path}")

    if failed_count:
        sys.exit(1)


if __name__ == "__main__":
    main()