"""Extract per-lane polylines from binary (JPEG-compressed) lane-segmentation masks.

The dataset at hand stores lanes as a single-class binary mask (white=lane,
black=background) rather than TuSimple's official per-lane (x,y) JSON annotations, so
lane *instances* have to be recovered from connected components of the thresholded
mask. Masks are lossy JPEG, so we clean small compression-noise blobs before labeling.
"""
from __future__ import annotations

import cv2
import numpy as np

MIN_BLOB_AREA = 150          # drop connected components smaller than this (JPEG noise)
BINARY_THRESHOLD = 127
MORPH_KERNEL = np.ones((3, 3), np.uint8)
MIN_POINTS_PER_LANE = 5       # discard components too short to be a real lane


def mask_to_lanes(mask: np.ndarray) -> list[list[tuple[float, float]]]:
    """Convert an (H,W) or (H,W,3) binary-ish mask into a list of lane polylines.

    Each polyline is a list of (x, y) pixel points, one point per mask row that the
    lane instance covers, ordered top-to-bottom (increasing y).
    """
    if mask.ndim == 3:
        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)

    binary = (mask > BINARY_THRESHOLD).astype(np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, MORPH_KERNEL)

    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)

    lanes: list[list[tuple[float, float]]] = []
    for label_id in range(1, num_labels):  # 0 is background
        area = stats[label_id, cv2.CC_STAT_AREA]
        if area < MIN_BLOB_AREA:
            continue

        ys, xs = np.where(labels == label_id)
        if ys.size == 0:
            continue

        row_to_xs: dict[int, list[int]] = {}
        for y, x in zip(ys, xs):
            row_to_xs.setdefault(int(y), []).append(int(x))

        points = [(float(np.mean(row_to_xs[y])), float(y)) for y in sorted(row_to_xs)]
        if len(points) >= MIN_POINTS_PER_LANE:
            lanes.append(points)

    return lanes