from __future__ import annotations

import argparse
import html
import os
from dataclasses import dataclass
from pathlib import Path

import cv2

os.environ.setdefault(
    "MPLCONFIGDIR",
    str(Path(__file__).resolve().parents[1] / ".matplotlib-cache"),
)

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}


@dataclass(frozen=True)
class TriangleConnection:
    row: int
    col: int
    corner: str


@dataclass(frozen=True)
class MatrixMapping:
    matrix: np.ndarray
    rendered_mask: np.ndarray
    x_edges: np.ndarray
    y_edges: np.ndarray
    triangles: tuple[TriangleConnection, ...] = ()
    grid_overlay: np.ndarray | None = None


def read_image(path: Path) -> np.ndarray:
    raw = np.fromfile(str(path), dtype=np.uint8)
    image = cv2.imdecode(raw, cv2.IMREAD_COLOR)
    if image is None:
        raise ValueError(f"Cannot read image: {path}")
    return image


def write_image(path: Path, image: np.ndarray) -> None:
    ok, encoded = cv2.imencode(path.suffix, image)
    if not ok:
        raise ValueError(f"Cannot encode image: {path}")
    encoded.tofile(str(path))


def binarize_foreground(gray: np.ndarray) -> np.ndarray:
    threshold, _ = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
    )
    gray_min = int(gray.min())
    gray_max = int(gray.max())
    if threshold <= gray_min or threshold >= gray_max:
        threshold = (gray_min + gray_max) / 2.0

    border = np.concatenate(
        [
            gray[0, :],
            gray[-1, :],
            gray[:, 0],
            gray[:, -1],
        ]
    )
    background_is_light = np.median(border) >= 128

    if background_is_light:
        foreground = (gray < threshold).astype(np.uint8) * 255
    else:
        foreground = (gray > threshold).astype(np.uint8) * 255

    return foreground


def fill_mask_holes(mask: np.ndarray) -> np.ndarray:
    foreground = mask > 0
    inverse = (~foreground).astype(np.uint8)
    flood = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)
    cv2.floodFill(inverse, flood, (0, 0), 2)
    holes = inverse == 1
    filled = foreground | holes
    return filled.astype(np.uint8) * 255


def fill_small_mask_holes(mask: np.ndarray, max_area: int) -> np.ndarray:
    if max_area <= 0:
        return mask

    foreground = mask > 0
    inverse = (~foreground).astype(np.uint8)
    flood = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)
    cv2.floodFill(inverse, flood, (0, 0), 2)
    holes = (flood == 1).astype(np.uint8)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
        holes, connectivity=8
    )

    filled = foreground.copy()
    for label in range(1, num_labels):
        area = int(stats[label, cv2.CC_STAT_AREA])
        if area <= max_area:
            filled[labels == label] = True

    return filled.astype(np.uint8) * 255


def fill_small_background_pockets(
    mask: np.ndarray,
    max_area: int,
    channel_size: int,
) -> np.ndarray:
    if max_area <= 0 or channel_size <= 1:
        return fill_small_mask_holes(mask, max_area)

    if channel_size % 2 == 0:
        channel_size += 1

    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (channel_size, channel_size)
    )
    expanded = cv2.dilate(mask, kernel)
    external = (expanded == 0).astype(np.uint8)
    flood = np.zeros((mask.shape[0] + 2, mask.shape[1] + 2), dtype=np.uint8)
    cv2.floodFill(external, flood, (0, 0), 2)

    pockets = ((mask == 0) & (external != 2)).astype(np.uint8)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
        pockets, connectivity=8
    )

    filled = mask > 0
    for label in range(1, num_labels):
        area = int(stats[label, cv2.CC_STAT_AREA])
        if area <= max_area:
            filled[labels == label] = True

    return filled.astype(np.uint8) * 255


def remove_small_components(mask: np.ndarray, min_area: int) -> np.ndarray:
    if min_area <= 0:
        return mask

    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
        mask, connectivity=8
    )
    cleaned = np.zeros_like(mask)
    for label in range(1, num_labels):
        if int(stats[label, cv2.CC_STAT_AREA]) >= min_area:
            cleaned[labels == label] = 255

    return cleaned


def smooth_binary_edges(mask: np.ndarray, iterations: int = 1) -> np.ndarray:
    foreground = mask > 0
    kernel = np.ones((3, 3), dtype=np.uint8)

    for _ in range(iterations):
        neighbors = cv2.filter2D(
            foreground.astype(np.uint8),
            ddepth=-1,
            kernel=kernel,
            borderType=cv2.BORDER_CONSTANT,
        )
        foreground = (foreground & (neighbors >= 3)) | (~foreground & (neighbors >= 7))

    return foreground.astype(np.uint8) * 255


def clean_mask(
    mask: np.ndarray,
    open_kernel_size: int,
    close_kernel_size: int,
    median_size: int,
    min_component_area: int,
    grid_size: int = 17,
) -> np.ndarray:
    if median_size > 1:
        if median_size % 2 == 0:
            median_size += 1
        mask = cv2.medianBlur(mask, median_size)

    cleaned = mask
    if open_kernel_size > 1:
        open_kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT, (open_kernel_size, open_kernel_size)
        )
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, open_kernel)

    cleaned = remove_small_components(cleaned, min_component_area)
    x1, y1, x2, y2 = _foreground_bbox(cleaned)
    cell_size = max(4.0, min(float(x2 - x1), float(y2 - y1)) / float(grid_size))
    burr_size = min(5, max(3, int(round(cell_size * 0.12)) | 1))
    line_length = max(3, int(round(cell_size * 0.64)))
    pocket_size = min(5, max(3, int(round(cell_size * 0.10)) | 1))
    gap_fill_length = min(7, max(3, int(round(cell_size * 0.18)) | 1))
    max_hole_area = max(4, int(round(cell_size * cell_size * 0.35)))

    burr_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (burr_size, burr_size)
    )
    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, burr_kernel)
    cleaned = remove_small_components(cleaned, min_component_area)

    horizontal_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (line_length, 1)
    )
    vertical_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (1, line_length)
    )
    horizontal = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, horizontal_kernel)
    vertical = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, vertical_kernel)
    cleaned = cv2.bitwise_or(horizontal, vertical)

    horizontal_gap_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (gap_fill_length, 1)
    )
    vertical_gap_kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT, (1, gap_fill_length)
    )
    horizontal_filled = cv2.morphologyEx(
        cleaned, cv2.MORPH_CLOSE, horizontal_gap_kernel
    )
    vertical_filled = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, vertical_gap_kernel)
    cleaned = cv2.bitwise_or(horizontal_filled, vertical_filled)

    cleaned = fill_small_background_pockets(
        cleaned,
        max_area=max_hole_area,
        channel_size=pocket_size,
    )
    smooth_size = close_kernel_size if close_kernel_size > 1 else 3
    smooth_kernel = cv2.getStructuringElement(
        cv2.MORPH_CROSS, (smooth_size, smooth_size)
    )
    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, smooth_kernel)
    cleaned = fill_small_background_pockets(
        cleaned,
        max_area=max_hole_area,
        channel_size=pocket_size,
    )
    cleaned = smooth_binary_edges(cleaned, iterations=1)

    return remove_small_components(cleaned, min_component_area)


def _render_matrix(
    matrix: np.ndarray,
    x_edges: np.ndarray,
    y_edges: np.ndarray,
    width: int,
    height: int,
    triangles: tuple[TriangleConnection, ...] = (),
) -> np.ndarray:
    rendered = np.zeros((height, width), dtype=np.uint8)

    rows, cols = matrix.shape
    for row_index in range(rows):
        y1 = int(y_edges[row_index])
        y2 = int(y_edges[row_index + 1])
        if y2 <= y1:
            continue

        for col_index in range(cols):
            if matrix[row_index, col_index] == 0:
                continue

            x1 = int(x_edges[col_index])
            x2 = int(x_edges[col_index + 1])
            if x2 > x1:
                rendered[y1:y2, x1:x2] = 255

    for triangle in triangles:
        x1 = int(x_edges[triangle.col])
        x2 = int(x_edges[triangle.col + 1])
        y1 = int(y_edges[triangle.row])
        y2 = int(y_edges[triangle.row + 1])
        points_by_corner = {
            "tl": [(x1, y1), (x2, y1), (x1, y2)],
            "tr": [(x2, y1), (x1, y1), (x2, y2)],
            "br": [(x2, y2), (x2, y1), (x1, y2)],
            "bl": [(x1, y2), (x1, y1), (x2, y2)],
        }
        points = np.array(points_by_corner[triangle.corner], dtype=np.int32)
        cv2.fillConvexPoly(rendered, points, 255)

    return rendered


def pixelate(mask: np.ndarray, factor: int) -> MatrixMapping:
    height, width = mask.shape
    small_width = max(1, round(width / factor))
    small_height = max(1, round(height / factor))
    sampled = cv2.resize(
        mask, (small_width, small_height), interpolation=cv2.INTER_AREA
    )
    small = np.where(sampled >= 96, 255, 0).astype(np.uint8)
    x_edges = np.rint(np.linspace(0, width, small_width + 1)).astype(int)
    y_edges = np.rint(np.linspace(0, height, small_height + 1)).astype(int)
    large = cv2.resize(small, (width, height), interpolation=cv2.INTER_NEAREST)
    return MatrixMapping(small, large, x_edges, y_edges)


def _order_quad_points(points: np.ndarray) -> np.ndarray:
    points = points.astype(np.float32)
    sums = points.sum(axis=1)
    diffs = np.diff(points, axis=1).reshape(-1)
    ordered = np.zeros((4, 2), dtype=np.float32)
    ordered[0] = points[int(np.argmin(sums))]
    ordered[2] = points[int(np.argmax(sums))]
    ordered[1] = points[int(np.argmin(diffs))]
    ordered[3] = points[int(np.argmax(diffs))]
    return ordered


def _foreground_quad(mask: np.ndarray) -> np.ndarray:
    points = cv2.findNonZero(mask)
    if points is None:
        height, width = mask.shape
        return np.array(
            [[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]],
            dtype=np.float32,
        )

    contour_points = points.reshape(-1, 2)
    hull = cv2.convexHull(contour_points)
    perimeter = cv2.arcLength(hull, True)
    approx = cv2.approxPolyDP(hull, 0.035 * perimeter, True)
    if len(approx) == 4:
        return _order_quad_points(approx.reshape(4, 2))

    rect = cv2.minAreaRect(contour_points)
    return _order_quad_points(cv2.boxPoints(rect))


def _quad_bbox_square_edges(
    quad: np.ndarray,
    canvas_width: int,
    canvas_height: int,
    grid_size: int,
) -> tuple[np.ndarray, np.ndarray]:
    x_min, y_min = np.min(quad, axis=0)
    x_max, y_max = np.max(quad, axis=0)
    center_x = float((x_min + x_max) * 0.5)
    center_y = float((y_min + y_max) * 0.5)
    side = max(float(x_max - x_min), float(y_max - y_min), 1.0)

    left = center_x - side * 0.5
    top = center_y - side * 0.5
    if left < 0:
        left = 0.0
    if top < 0:
        top = 0.0
    if left + side > canvas_width:
        left = max(0.0, float(canvas_width) - side)
    if top + side > canvas_height:
        top = max(0.0, float(canvas_height) - side)

    right = min(float(canvas_width), left + side)
    bottom = min(float(canvas_height), top + side)
    x_edges = np.rint(np.linspace(left, right, grid_size + 1)).astype(int)
    y_edges = np.rint(np.linspace(top, bottom, grid_size + 1)).astype(int)
    return x_edges, y_edges


def _warp_to_grid(mask: np.ndarray, grid_size: int, cell_samples: int) -> np.ndarray:
    quad = _foreground_quad(mask)
    side = int(grid_size * cell_samples)
    destination = np.array(
        [[0, 0], [side - 1, 0], [side - 1, side - 1], [0, side - 1]],
        dtype=np.float32,
    )
    transform = cv2.getPerspectiveTransform(quad, destination)
    return cv2.warpPerspective(
        mask,
        transform,
        (side, side),
        flags=cv2.INTER_NEAREST,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )


def _foreground_bbox(mask: np.ndarray) -> tuple[int, int, int, int]:
    ys, xs = np.where(mask > 0)
    if xs.size == 0:
        height, width = mask.shape
        return 0, 0, width, height

    return int(xs.min()), int(ys.min()), int(xs.max() + 1), int(ys.max() + 1)


def _rotate_mask(mask: np.ndarray, angle: float) -> np.ndarray:
    if abs(angle) < 1e-6:
        return mask

    height, width = mask.shape
    center = (width / 2.0, height / 2.0)
    transform = cv2.getRotationMatrix2D(center, angle, 1.0)
    return cv2.warpAffine(
        mask,
        transform,
        (width, height),
        flags=cv2.INTER_NEAREST,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )


def _projection_score(mask: np.ndarray) -> float:
    foreground = mask > 0
    row_sums = foreground.sum(axis=1).astype(np.float64)
    col_sums = foreground.sum(axis=0).astype(np.float64)
    return float(np.square(row_sums).sum() + np.square(col_sums).sum())


def _estimate_skew_angle(mask: np.ndarray, max_angle: float = 4.0) -> float:
    best_angle = 0.0
    best_score = _projection_score(mask)

    for angle in np.arange(-max_angle, max_angle + 0.001, 0.25):
        if abs(float(angle)) < 1e-6:
            continue
        rotated = _rotate_mask(mask, float(angle))
        score = _projection_score(rotated)
        if score > best_score:
            best_score = score
            best_angle = float(angle)

    return best_angle


def _solidify_hollow_strokes(mask: np.ndarray, grid_size: int) -> np.ndarray:
    x1, y1, x2, y2 = _foreground_bbox(mask)
    cell_width = max(2.0, float(x2 - x1) / float(grid_size))
    cell_height = max(2.0, float(y2 - y1) / float(grid_size))
    cell_size = max(2.0, min(cell_width, cell_height))

    long_size = max(3, int(round(cell_size * 0.72)) | 1)
    short_size = max(3, int(round(cell_size * 0.14)) | 1)

    vertical_bridge = cv2.getStructuringElement(
        cv2.MORPH_RECT, (short_size, long_size)
    )
    horizontal_bridge = cv2.getStructuringElement(
        cv2.MORPH_RECT, (long_size, short_size)
    )
    square_cleanup = cv2.getStructuringElement(
        cv2.MORPH_RECT, (max(3, short_size), max(3, short_size))
    )

    horizontal_strokes = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, vertical_bridge)
    vertical_strokes = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, horizontal_bridge)
    solid = cv2.bitwise_or(horizontal_strokes, vertical_strokes)
    solid = cv2.morphologyEx(solid, cv2.MORPH_CLOSE, square_cleanup)
    return solid


def _triangle_corner_score(cell: np.ndarray, corner: str) -> float:
    size = cell.shape[0]
    yy, xx = np.indices(cell.shape)
    if corner == "tl":
        region = (xx + yy) <= size
    elif corner == "tr":
        region = ((size - 1 - xx) + yy) <= size
    elif corner == "br":
        region = ((size - 1 - xx) + (size - 1 - yy)) <= size
    else:
        region = (xx + (size - 1 - yy)) <= size

    foreground = cell > 0
    inside = float(foreground[region].mean()) if np.any(region) else 0.0
    outside = float(foreground[~region].mean()) if np.any(~region) else 0.0
    return inside - outside


def _detect_triangle_connections(
    warped: np.ndarray,
    matrix: np.ndarray,
    cell_samples: int,
    triangle_threshold: float,
) -> tuple[TriangleConnection, ...]:
    grid_size = matrix.shape[0]
    triangles: list[TriangleConnection] = []

    def add_if_supported(row: int, col: int, corner: str) -> None:
        y1 = row * cell_samples
        y2 = y1 + cell_samples
        x1 = col * cell_samples
        x2 = x1 + cell_samples
        cell = warped[y1:y2, x1:x2]
        coverage = float((cell > 0).mean())
        score = _triangle_corner_score(cell, corner)
        if coverage >= triangle_threshold and score >= triangle_threshold:
            triangles.append(TriangleConnection(row, col, corner))

    for row in range(grid_size - 1):
        for col in range(grid_size - 1):
            nw = matrix[row, col] > 0
            ne = matrix[row, col + 1] > 0
            sw = matrix[row + 1, col] > 0
            se = matrix[row + 1, col + 1] > 0

            if nw and se and not ne and not sw:
                add_if_supported(row, col + 1, "bl")
                add_if_supported(row + 1, col, "tr")
            elif ne and sw and not nw and not se:
                add_if_supported(row, col, "br")
                add_if_supported(row + 1, col + 1, "tl")

    return tuple(triangles)


def _normalize_raw_strokes(
    mask: np.ndarray,
    grid_size: int,
    cell_samples: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    angle = _estimate_skew_angle(mask)
    deskewed = _rotate_mask(mask, angle)
    x1, y1, x2, y2 = _foreground_bbox(deskewed)
    side = int(grid_size * cell_samples)
    candidates: list[tuple[float, np.ndarray, tuple[float, ...]]] = []
    scale_values = (0.92, 0.94, 0.96, 0.98, 1.0, 1.02, 1.04, 1.06)
    shift_x_values = (
        -0.5,
        -0.4,
        -0.3,
        -0.2,
        -0.1,
        0.0,
        0.1,
        0.2,
        0.3,
        0.4,
        0.5,
    )
    shift_y_values = (-0.3, -0.2, -0.1, 0.0, 0.1, 0.2, 0.3)
    for scale_x in scale_values:
        for scale_y in scale_values:
            for shift_x in shift_x_values:
                for shift_y in shift_y_values:
                    warped = _warp_bbox_candidate(
                        deskewed,
                        x1,
                        y1,
                        x2,
                        y2,
                        scale_x,
                        scale_y,
                        shift_x,
                        shift_y,
                        side,
                    )
                    candidates.append(
                        (
                            _grid_fit_score(warped, grid_size),
                            warped,
                            (
                                scale_x,
                                scale_y,
                                shift_x,
                                shift_y,
                                0.0,
                                0.0,
                                0.0,
                                0.0,
                            ),
                        )
                    )

    for _score, _warped, params in sorted(
        candidates,
        reverse=True,
        key=lambda item: item[0],
    )[:36]:
        (
            scale_x,
            scale_y,
            shift_x,
            shift_y,
            _shear_x,
            _shear_y,
            _curve_x,
            _curve_y,
        ) = params
        for shear_x in (-0.8, -0.4, 0.0, 0.4, 0.8):
            for shear_y in (-0.6, -0.3, 0.0, 0.3, 0.6):
                warped = _warp_bbox_candidate(
                    deskewed,
                    x1,
                    y1,
                    x2,
                    y2,
                    scale_x,
                    scale_y,
                    shift_x,
                    shift_y,
                    side,
                    shear_x=shear_x,
                    shear_y=shear_y,
                )
                candidates.append(
                    (
                        _grid_fit_score(warped, grid_size),
                        warped,
                        (
                            scale_x,
                            scale_y,
                            shift_x,
                            shift_y,
                            shear_x,
                            shear_y,
                            0.0,
                            0.0,
                        ),
                    )
                )

    for _score, _warped, params in sorted(
        candidates,
        reverse=True,
        key=lambda item: item[0],
    )[:28]:
        (
            scale_x,
            scale_y,
            shift_x,
            shift_y,
            shear_x,
            shear_y,
            _curve_x,
            _curve_y,
        ) = params
        for curve_x, curve_y in (
            (-0.65, 0.0),
            (-0.35, 0.0),
            (0.35, 0.0),
            (0.65, 0.0),
            (0.0, -0.65),
            (0.0, -0.35),
            (0.0, 0.35),
            (0.0, 0.65),
            (-0.35, -0.35),
            (-0.35, 0.35),
            (0.35, -0.35),
            (0.35, 0.35),
        ):
            warped = _warp_bbox_candidate(
                deskewed,
                x1,
                y1,
                x2,
                y2,
                scale_x,
                scale_y,
                shift_x,
                shift_y,
                side,
                shear_x=shear_x,
                shear_y=shear_y,
                curve_x=curve_x,
                curve_y=curve_y,
            )
            candidates.append(
                (
                    _grid_fit_score(warped, grid_size),
                    warped,
                    (
                        scale_x,
                        scale_y,
                        shift_x,
                        shift_y,
                        shear_x,
                        shear_y,
                        curve_x,
                        curve_y,
                    ),
                )
            )

    if candidates:
        warped = max(candidates, key=lambda item: item[0])[1]
    else:
        crop = deskewed[y1:y2, x1:x2]
        if crop.size == 0:
            crop = deskewed
        warped = cv2.resize(crop, (side, side), interpolation=cv2.INTER_NEAREST)

    render_x1, render_y1, render_x2, render_y2 = _foreground_bbox(mask)
    x_edges = np.rint(np.linspace(render_x1, render_x2, grid_size + 1)).astype(int)
    y_edges = np.rint(np.linspace(render_y1, render_y2, grid_size + 1)).astype(int)
    return warped, x_edges, y_edges, _grid_overlay_preview(warped, grid_size)


def _grid_overlay_preview(warped: np.ndarray, grid_size: int) -> np.ndarray:
    preview = cv2.cvtColor(mask_to_display(warped), cv2.COLOR_GRAY2RGB)
    height, width = warped.shape
    for index in range(grid_size + 1):
        x = int(round(index * width / grid_size))
        y = int(round(index * height / grid_size))
        thickness = 2 if index in (0, grid_size) else 1
        cv2.line(preview, (x, 0), (x, height - 1), (220, 40, 40), thickness)
        cv2.line(preview, (0, y), (width - 1, y), (220, 40, 40), thickness)
    return preview


def _warp_bbox_candidate(
    mask: np.ndarray,
    x1: int,
    y1: int,
    x2: int,
    y2: int,
    scale_x: float,
    scale_y: float,
    shift_x: float,
    shift_y: float,
    side: int,
    shear_x: float = 0.0,
    shear_y: float = 0.0,
    curve_x: float = 0.0,
    curve_y: float = 0.0,
) -> np.ndarray:
    width = max(1.0, float(x2 - x1))
    height = max(1.0, float(y2 - y1))
    center_x = ((x1 + x2) * 0.5) + shift_x * width / 17.0
    center_y = ((y1 + y2) * 0.5) + shift_y * height / 17.0
    crop_width = width * scale_x
    crop_height = height * scale_y
    left = int(round(center_x - crop_width * 0.5))
    right = int(round(center_x + crop_width * 0.5))
    top = int(round(center_y - crop_height * 0.5))
    bottom = int(round(center_y + crop_height * 0.5))

    crop = np.zeros((max(1, bottom - top), max(1, right - left)), dtype=np.uint8)
    source_left = max(0, left)
    source_right = min(mask.shape[1], right)
    source_top = max(0, top)
    source_bottom = min(mask.shape[0], bottom)
    if source_right > source_left and source_bottom > source_top:
        crop[
            source_top - top : source_bottom - top,
            source_left - left : source_right - left,
        ] = mask[source_top:source_bottom, source_left:source_right]

    if (
        abs(shear_x) < 1e-9
        and abs(shear_y) < 1e-9
        and abs(curve_x) < 1e-9
        and abs(curve_y) < 1e-9
    ):
        return cv2.resize(crop, (side, side), interpolation=cv2.INTER_NEAREST)

    yy, xx = np.indices((side, side), dtype=np.float32)
    u = xx / max(1.0, float(side - 1))
    v = yy / max(1.0, float(side - 1))
    centered_u = u - 0.5
    centered_v = v - 0.5
    map_x = (
        u * float(crop.shape[1] - 1)
        + (shear_x * centered_v + curve_x * (centered_v * centered_v - 0.125))
        * float(crop.shape[1])
        / 17.0
    )
    map_y = (
        v * float(crop.shape[0] - 1)
        + (shear_y * centered_u + curve_y * (centered_u * centered_u - 0.125))
        * float(crop.shape[0])
        / 17.0
    )
    return cv2.remap(
        crop,
        map_x,
        map_y,
        interpolation=cv2.INTER_NEAREST,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=0,
    )


def _grid_fit_score(warped: np.ndarray, grid_size: int) -> float:
    occupancy = cv2.resize(
        (warped > 0).astype(np.float32),
        (grid_size, grid_size),
        interpolation=cv2.INTER_AREA,
    )
    polar_score = float((np.abs(occupancy - 0.5) * 2.0).mean())
    mixed_cells = float(((occupancy > 0.08) & (occupancy < 0.70)).mean())
    strong_cells = float((occupancy > 0.20).mean())
    foreground_density = float(occupancy.mean())
    border_density = float(
        (
            occupancy[0].mean()
            + occupancy[-1].mean()
            + occupancy[:, 0].mean()
            + occupancy[:, -1].mean()
        )
        * 0.25
    )
    expected_strong_cells = 0.76
    expected_density = 0.58
    return (
        polar_score
        - 0.18 * mixed_cells
        - 0.35 * abs(strong_cells - expected_strong_cells)
        - 0.20 * abs(foreground_density - expected_density)
        + 0.25 * border_density
    )


def _cell_ratio_features(
    warped: np.ndarray,
    grid_size: int,
    cell_samples: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    occupancy = np.zeros((grid_size, grid_size), dtype=np.float64)
    center = np.zeros((grid_size, grid_size), dtype=np.float64)
    sides = np.zeros((grid_size, grid_size, 4), dtype=np.float64)
    center_margin = max(1, cell_samples // 4)
    side_width = max(1, cell_samples // 3)

    for row in range(grid_size):
        for col in range(grid_size):
            y1 = row * cell_samples
            y2 = y1 + cell_samples
            x1 = col * cell_samples
            x2 = x1 + cell_samples
            cell = (warped[y1:y2, x1:x2] > 0).astype(np.float32)
            occupancy[row, col] = float(cell.mean())
            center_cell = cell[
                center_margin : cell_samples - center_margin,
                center_margin : cell_samples - center_margin,
            ]
            center[row, col] = float(center_cell.mean()) if center_cell.size else 0.0
            sides[row, col] = (
                float(cell[:, :side_width].mean()),
                float(cell[:, -side_width:].mean()),
                float(cell[:side_width, :].mean()),
                float(cell[-side_width:, :].mean()),
            )

    return occupancy, center, sides


def _ratio_matrix(
    warped: np.ndarray,
    grid_size: int,
    cell_samples: int,
    threshold: float = 0.44,
) -> np.ndarray:
    occupancy, center, sides = _cell_ratio_features(
        warped,
        grid_size=grid_size,
        cell_samples=cell_samples,
    )
    score = 0.65 * occupancy + 0.35 * center
    matrix = (score >= threshold).astype(np.uint8)
    matrix = _resolve_thick_stroke_overlaps(matrix, score, sides)
    return matrix.astype(np.uint8) * 255


def _resolve_thick_stroke_overlaps(
    matrix: np.ndarray,
    score: np.ndarray,
    sides: np.ndarray,
) -> np.ndarray:
    resolved = matrix.astype(bool).copy()
    rows, cols = resolved.shape

    def active_neighbors(row: int, col: int) -> int:
        total = 0
        for n_row, n_col in (
            (row - 1, col),
            (row + 1, col),
            (row, col - 1),
            (row, col + 1),
        ):
            if 0 <= n_row < rows and 0 <= n_col < cols and resolved[n_row, n_col]:
                total += 1
        return total

    def bridge_score(row: int, col: int) -> int:
        bridge = 0
        if col > 0 and col + 1 < cols and resolved[row, col - 1] and resolved[row, col + 1]:
            bridge += 1
        if row > 0 and row + 1 < rows and resolved[row - 1, col] and resolved[row + 1, col]:
            bridge += 1
        return bridge

    for _iteration in range(rows * cols):
        weakest: tuple[float, int, int] | None = None
        for row in range(rows - 1):
            for col in range(cols - 1):
                if int(resolved[row : row + 2, col : col + 2].sum()) != 4:
                    continue

                for cell_row in (row, row + 1):
                    for cell_col in (col, col + 1):
                        side_imbalance = float(
                            sides[cell_row, cell_col].max()
                            - sides[cell_row, cell_col].min()
                        )
                        support = (
                            float(score[cell_row, cell_col])
                            + 0.10 * active_neighbors(cell_row, cell_col)
                            + 0.25 * bridge_score(cell_row, cell_col)
                            - 0.05 * side_imbalance
                        )
                        candidate = (support, cell_row, cell_col)
                        if weakest is None or candidate[0] < weakest[0]:
                            weakest = candidate

        if weakest is None:
            break
        _, row, col = weakest
        resolved[row, col] = False

    return resolved.astype(np.uint8)


def count_2x2_blocks(matrix: np.ndarray) -> int:
    image = (matrix > 0).astype(np.uint8)
    total = 0
    for row in range(image.shape[0] - 1):
        for col in range(image.shape[1] - 1):
            total += int(image[row : row + 2, col : col + 2].sum() == 4)
    return total


def matrix_map(
    mask: np.ndarray,
    grid_size: int,
    triangle_threshold: float,
    cell_samples: int,
) -> MatrixMapping:
    height, width = mask.shape
    warped, x_edges, y_edges, grid_overlay = _normalize_raw_strokes(
        mask,
        grid_size=grid_size,
        cell_samples=cell_samples,
    )
    matrix = _ratio_matrix(
        warped,
        grid_size=grid_size,
        cell_samples=cell_samples,
    )
    triangles = _detect_triangle_connections(
        warped,
        matrix,
        cell_samples=cell_samples,
        triangle_threshold=triangle_threshold,
    )

    rendered = _render_matrix(matrix, x_edges, y_edges, width, height, triangles)
    return MatrixMapping(
        matrix,
        rendered,
        x_edges,
        y_edges,
        triangles=triangles,
        grid_overlay=grid_overlay,
    )


def mask_to_display(mask: np.ndarray) -> np.ndarray:
    return np.where(mask > 0, 0, 255).astype(np.uint8)


def save_comparison(
    path: Path,
    original_bgr: np.ndarray,
    cleaned_mask: np.ndarray,
    grid_overlay: np.ndarray | None,
    mapped_mask: np.ndarray,
    mapped_title: str,
) -> None:
    original_rgb = cv2.cvtColor(original_bgr, cv2.COLOR_BGR2RGB)
    cleaned = mask_to_display(cleaned_mask)
    mapped = mask_to_display(mapped_mask)

    fig, axes = plt.subplots(1, 4, figsize=(16, 4), constrained_layout=True)
    panels = [
        ("Original", original_rgb, None),
        ("Cleaned binary", cleaned, "gray"),
        ("Adapted 17x17 grid", grid_overlay if grid_overlay is not None else cleaned, None),
        (mapped_title, mapped, "gray"),
    ]

    for axis, (title, image, cmap) in zip(axes, panels):
        axis.imshow(image, cmap=cmap)
        axis.set_title(title)
        axis.axis("off")

    fig.savefig(path, dpi=180)
    plt.close(fig)


def small_matrix_to_svg(
    small_mask: np.ndarray,
    original_width: int,
    original_height: int,
    title: str,
    x_edges: np.ndarray,
    y_edges: np.ndarray,
    triangles: tuple[TriangleConnection, ...],
) -> str:
    rows, cols = small_mask.shape
    escaped_title = html.escape(title)
    parts = [
        '<?xml version="1.0" encoding="UTF-8"?>',
        (
            f'<svg xmlns="http://www.w3.org/2000/svg" '
            f'width="{original_width}" height="{original_height}" '
            f'viewBox="0 0 {original_width} {original_height}" '
            f'shape-rendering="crispEdges">'
        ),
        f"  <title>{escaped_title}</title>",
        '  <rect width="100%" height="100%" fill="white"/>',
        '  <g fill="black">',
    ]

    for row_index in range(rows):
        col_index = 0
        while col_index < cols:
            if small_mask[row_index, col_index] == 0:
                col_index += 1
                continue

            run_start = col_index
            while col_index < cols and small_mask[row_index, col_index] > 0:
                col_index += 1

            x = int(x_edges[run_start])
            y = int(y_edges[row_index])
            width = int(x_edges[col_index]) - x
            height = int(y_edges[row_index + 1]) - y
            parts.append(
                f'    <rect x="{x:.4f}" y="{y:.4f}" '
                f'width="{width:.4f}" height="{height:.4f}"/>'
            )

    for triangle in triangles:
        x1 = int(x_edges[triangle.col])
        x2 = int(x_edges[triangle.col + 1])
        y1 = int(y_edges[triangle.row])
        y2 = int(y_edges[triangle.row + 1])
        points_by_corner = {
            "tl": [(x1, y1), (x2, y1), (x1, y2)],
            "tr": [(x2, y1), (x1, y1), (x2, y2)],
            "br": [(x2, y2), (x2, y1), (x1, y2)],
            "bl": [(x1, y2), (x1, y1), (x2, y2)],
        }
        points = " ".join(
            f"{point_x:.4f},{point_y:.4f}"
            for point_x, point_y in points_by_corner[triangle.corner]
        )
        parts.append(f'    <polygon points="{points}"/>')

    parts.extend(["  </g>", "</svg>", ""])
    return "\n".join(parts)


def matrix_to_text(
    matrix: np.ndarray,
    triangles: tuple[TriangleConnection, ...],
) -> str:
    lines = []
    for row in matrix:
        lines.append("".join("1" if value > 0 else "0" for value in row))

    lines.append("")
    lines.append(f"# 2x2_black_blocks: {count_2x2_blocks(matrix)}")

    if triangles:
        lines.append("")
        lines.append("# triangles: row,col,corner")
        for triangle in triangles:
            lines.append(f"{triangle.row},{triangle.col},{triangle.corner}")

    return "\n".join(lines) + "\n"


def process_image(
    image_path: Path,
    processed_dir: Path,
    comparison_dir: Path,
    svg_dir: Path,
    factor: int,
    grid_size: int,
    cell_samples: int,
    open_kernel_size: int,
    close_kernel_size: int,
    median_size: int,
    min_component_area: int,
    mode: str,
    triangle_threshold: float,
) -> int:
    original = read_image(image_path)
    gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
    binary = binarize_foreground(gray)
    cleaned = clean_mask(
        binary,
        open_kernel_size=open_kernel_size,
        close_kernel_size=close_kernel_size,
        median_size=median_size,
        min_component_area=min_component_area,
        grid_size=grid_size,
    )
    if mode == "matrix":
        mapping = matrix_map(
            cleaned,
            grid_size=grid_size,
            triangle_threshold=triangle_threshold,
            cell_samples=cell_samples,
        )
        mapped_title = f"{grid_size}x{grid_size} matrix"
    else:
        mapping = pixelate(cleaned, factor=factor)
        mapped_title = "Pixelated"

    stem = image_path.stem
    write_image(processed_dir / f"{stem}_cleaned.png", mask_to_display(cleaned))
    write_image(
        processed_dir / f"{stem}_pixelated.png", mask_to_display(mapping.rendered_mask)
    )
    write_image(
        processed_dir / f"{stem}_matrix{mapping.matrix.shape[0]}.png",
        mask_to_display(mapping.matrix),
    )
    (processed_dir / f"{stem}_matrix{mapping.matrix.shape[0]}.txt").write_text(
        matrix_to_text(mapping.matrix, mapping.triangles),
        encoding="utf-8",
    )
    save_comparison(
        comparison_dir / f"{stem}_comparison.png",
        original,
        cleaned,
        mapping.grid_overlay,
        mapping.rendered_mask,
        mapped_title,
    )

    height, width = gray.shape
    svg = small_matrix_to_svg(
        mapping.matrix,
        width,
        height,
        stem,
        mapping.x_edges,
        mapping.y_edges,
        mapping.triangles,
    )
    (svg_dir / f"{stem}_pixelated.svg").write_text(svg, encoding="utf-8")
    return count_2x2_blocks(mapping.matrix)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Clean, map, and vectorize scanned seal-script characters."
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        default=Path("PreProcessing/original_characters"),
    )
    parser.add_argument(
        "--processed-dir",
        type=Path,
        default=Path("PreProcessing/processed_characters"),
    )
    parser.add_argument(
        "--comparison-dir",
        type=Path,
        default=Path("PreProcessing/comparisons"),
    )
    parser.add_argument(
        "--svg-dir",
        type=Path,
        default=Path("PreProcessing/svg_characters"),
    )
    parser.add_argument(
        "--factor",
        type=int,
        default=15,
        help="Only used by --mode resize.",
    )
    parser.add_argument(
        "--grid-size",
        type=int,
        default=17,
        help="Strict square matrix size used by --mode matrix.",
    )
    parser.add_argument(
        "--cell-samples",
        type=int,
        default=32,
        help="Samples per grid cell while deskewing and classifying the source.",
    )
    parser.add_argument(
        "--mode",
        choices=("matrix", "resize"),
        default="matrix",
        help="matrix maps to a strict grid; resize keeps the old mosaic-style downsample.",
    )
    parser.add_argument(
        "--triangle-threshold",
        type=float,
        default=0.18,
        help="Minimum triangular evidence in an empty diagonal-connection cell.",
    )
    parser.add_argument("--open-kernel-size", type=int, default=1)
    parser.add_argument(
        "--close-kernel-size",
        type=int,
        default=0,
        help="Optional fixed close kernel for clean; 0 uses a grid-relative size.",
    )
    parser.add_argument("--median-size", type=int, default=1)
    parser.add_argument("--min-component-area", type=int, default=8)
    args = parser.parse_args()

    for directory in (args.processed_dir, args.comparison_dir, args.svg_dir):
        directory.mkdir(parents=True, exist_ok=True)

    images = sorted(
        path
        for path in args.input_dir.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )
    if not images:
        raise SystemExit(f"No images found in {args.input_dir}")

    for image_path in images:
        two_by_two_blocks = process_image(
            image_path=image_path,
            processed_dir=args.processed_dir,
            comparison_dir=args.comparison_dir,
            svg_dir=args.svg_dir,
            factor=args.factor,
            grid_size=args.grid_size,
            cell_samples=args.cell_samples,
            open_kernel_size=args.open_kernel_size,
            close_kernel_size=args.close_kernel_size,
            median_size=args.median_size,
            min_component_area=args.min_component_area,
            mode=args.mode,
            triangle_threshold=args.triangle_threshold,
        )
        print(f"processed: {image_path.name} (2x2_black_blocks={two_by_two_blocks})")


if __name__ == "__main__":
    main()