#!/usr/bin/env python3
"""Generate an AprilTag border around a rectangle and save a printable image."""

import argparse
import csv
import math
import os
from pathlib import Path
import sys
from typing import List, Tuple

try:
    import numpy as np
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'numpy'. Install it with: pip install numpy"
    ) from exc

try:
    import cv2
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'opencv-contrib-python'. Install it with: pip install opencv-contrib-python"
    ) from exc

# Force non-GUI backend to avoid Qt/xcb plugin issues in headless environments.
os.environ.setdefault("MPLBACKEND", "Agg")
try:
    import matplotlib
    matplotlib.use("Agg", force=True)
    import matplotlib.pyplot as plt
except ImportError as exc:
    raise SystemExit(
        "Missing dependency 'matplotlib'. Install it with: pip install matplotlib"
    ) from exc


APRILTAG_FAMILIES = {
    "tag16h5": ("DICT_APRILTAG_16h5", 30),
    "tag25h9": ("DICT_APRILTAG_25h9", 35),
    "tag36h10": ("DICT_APRILTAG_36h10", 2320),
    "tag36h11": ("DICT_APRILTAG_36h11", 587),
}


A4_WIDTH_CM = 29.7
A4_HEIGHT_CM = 21.0
# A4_HEIGHT_CM = 29.7


def cm_to_px(cm: float, dpi: int) -> int:
    return int(round(cm / 2.54 * dpi))


def parse_ratio_text(value: str) -> float:
    text = value.strip()
    if not text:
        raise ValueError("empty ratio")
    if ":" in text:
        parts = text.split(":")
        if len(parts) != 2:
            raise ValueError(f"invalid ratio '{value}', expected like 1:4")
        left = float(parts[0].strip())
        right = float(parts[1].strip())
        if left < 0 or right <= 0:
            raise ValueError(f"invalid ratio '{value}', left must be >= 0 and right > 0")
        return left / right
    ratio = float(text)
    if ratio < 0:
        raise ValueError(f"invalid ratio '{value}', must be >= 0")
    return ratio


def border_positions(total_cm: float, tag_cm: float, pitch_cm: float) -> List[float]:
    if total_cm + 1e-9 < tag_cm:
        raise ValueError("rectangle side is smaller than one tag")

    count = int(math.floor((total_cm - tag_cm) / pitch_cm)) + 1
    count = max(1, count)
    used_cm = tag_cm + (count - 1) * pitch_cm
    margin_cm = (total_cm - used_cm) / 2.0
    return [margin_cm + i * pitch_cm for i in range(count)]


def side_positions_excluding_corners(total_cm: float, tag_cm: float, pitch_cm: float) -> List[float]:
    inner_cm = total_cm - 2.0 * tag_cm
    if inner_cm + 1e-9 < tag_cm:
        return []

    count = int(math.floor((inner_cm - tag_cm) / pitch_cm)) + 1
    count = max(1, count)
    used_cm = tag_cm + (count - 1) * pitch_cm
    margin_cm = (inner_cm - used_cm) / 2.0
    start_cm = tag_cm + margin_cm
    return [start_cm + i * pitch_cm for i in range(count)]


def resolve_aruco_dict(family: str) -> Tuple["cv2.aruco_Dictionary", int]:
    if family not in APRILTAG_FAMILIES:
        supported = ", ".join(sorted(APRILTAG_FAMILIES))
        raise ValueError(f"Unsupported family '{family}'. Supported: {supported}")

    dict_name, max_tags = APRILTAG_FAMILIES[family]
    dict_id = getattr(cv2.aruco, dict_name, None)
    if dict_id is None:
        raise RuntimeError(
            f"OpenCV build does not include {dict_name}. "
            "Please install opencv-contrib-python."
        )

    return cv2.aruco.getPredefinedDictionary(dict_id), max_tags


def make_marker(aruco_dict: "cv2.aruco_Dictionary", marker_id: int, size_px: int) -> np.ndarray:
    if hasattr(cv2.aruco, "generateImageMarker"):
        marker = cv2.aruco.generateImageMarker(aruco_dict, marker_id, size_px)
        if marker is not None:
            return marker

    marker = np.zeros((size_px, size_px), dtype=np.uint8)
    cv2.aruco.drawMarker(aruco_dict, marker_id, size_px, marker, 1)
    return marker


def save_printable_png(output_path: Path, canvas: np.ndarray, width_cm: float, height_cm: float, dpi: int) -> None:
    fig = plt.figure(figsize=(width_cm / 2.54, height_cm / 2.54), dpi=dpi, frameon=False)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(canvas, cmap="gray", vmin=0, vmax=255, interpolation="nearest")
    ax.axis("off")
    fig.savefig(output_path, dpi=dpi, pad_inches=0)
    plt.close(fig)


def save_a4_pdf(
    output_path: Path,
    canvas: np.ndarray,
    width_cm: float,
    height_cm: float,
    dpi: int,
    landscape: bool,
) -> None:
    if landscape:
        page_w_cm = A4_HEIGHT_CM
        page_h_cm = A4_WIDTH_CM
    else:
        page_w_cm = A4_WIDTH_CM
        page_h_cm = A4_HEIGHT_CM

    if width_cm > page_w_cm + 1e-9 or height_cm > page_h_cm + 1e-9:
        raise SystemExit(
            f"Pattern size {width_cm}x{height_cm} cm does not fit A4 "
            f"({'landscape' if landscape else 'portrait'}) {page_w_cm}x{page_h_cm} cm."
        )

    fig = plt.figure(figsize=(page_w_cm / 2.54, page_h_cm / 2.54), dpi=dpi, frameon=False)
    x0 = (page_w_cm - width_cm) / 2.0 / page_w_cm
    y0 = (page_h_cm - height_cm) / 2.0 / page_h_cm
    w = width_cm / page_w_cm
    h = height_cm / page_h_cm
    ax = fig.add_axes([x0, y0, w, h])
    ax.imshow(canvas, cmap="gray", vmin=0, vmax=255, interpolation="nearest")
    ax.axis("off")
    fig.savefig(output_path, format="pdf", dpi=dpi, pad_inches=0)
    plt.close(fig)


def write_layout_csv(path: Path, rows: List[Tuple[int, float, float, int, int]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["tag_id", "x_cm", "y_cm", "x_px", "y_px"])
        for row in rows:
            writer.writerow(row)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate AprilTag border image")
    parser.add_argument("--width-cm", type=float, default=26.0, help="Rectangle width in cm")
    parser.add_argument("--height-cm", type=float, default=17.0, help="Rectangle height in cm")
    parser.add_argument(
        "--tag-size-cm",
        type=float,
        default=2.0,
        help="AprilTag marker side length in cm (excluding per-tag white border)",
    )
    parser.add_argument(
        "--pitch-cm",
        type=float,
        default=None,
        help="Distance between adjacent tag tile top-left positions in cm (default: auto = tag tile size)",
    )
    parser.add_argument(
        "--family",
        type=str,
        default="tag36h11",
        choices=sorted(APRILTAG_FAMILIES.keys()),
        help="AprilTag family",
    )
    parser.add_argument("--start-id", type=int, default=0, help="First tag ID")
    parser.add_argument(
        "--wrap-ids",
        action="store_true",
        help="Wrap around if IDs exceed family max",
    )
    parser.add_argument(
        "--white-border-cm",
        type=float,
        default=0.0,
        help="Keep this much white margin on all four sides (cm)",
    )
    parser.add_argument(
        "--tag-white-ratio",
        type=str,
        default="1:4",
        help="Total white border to marker ratio (e.g. 1:4 means total white width is 1/4 of marker size)",
    )
    parser.add_argument(
        "--tag-white-border-cm",
        type=float,
        default=None,
        help="Optional override: white margin on each side in cm; if omitted, auto-computed from --tag-white-ratio",
    )
    parser.add_argument("--dpi", type=int, default=300, help="Output DPI")
    parser.add_argument(
        "--a4-landscape",
        action="store_true",
        help="When output is PDF, use A4 landscape instead of portrait",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("apriltags_border_20x15cm_A4.pdf"),
        help="Output path (.pdf for A4 page, .png for exact-size image)",
    )
    parser.add_argument(
        "--layout-csv",
        type=Path,
        default=None,
        help="Optional CSV path to save tag placements",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    if args.width_cm <= 0 or args.height_cm <= 0 or args.tag_size_cm <= 0:
        raise SystemExit("width/height/tag-size must all be > 0")
    if args.white_border_cm < 0:
        raise SystemExit("white-border-cm must be >= 0")

    if args.tag_white_border_cm is None:
        try:
            tag_white_ratio = parse_ratio_text(args.tag_white_ratio)
        except ValueError as exc:
            raise SystemExit(f"tag-white-ratio parse error: {exc}") from exc
        tag_white_border_cm = args.tag_size_cm * tag_white_ratio / 2.0
    else:
        if args.tag_white_border_cm < 0:
            raise SystemExit("tag-white-border-cm must be >= 0")
        tag_white_border_cm = args.tag_white_border_cm
        tag_white_ratio = 2.0 * tag_white_border_cm / args.tag_size_cm

    tag_tile_size_cm = args.tag_size_cm + 2.0 * tag_white_border_cm
    pitch_cm = tag_tile_size_cm if args.pitch_cm is None else args.pitch_cm
    if pitch_cm <= 0:
        raise SystemExit("pitch-cm must be > 0")
    if pitch_cm + 1e-9 < tag_tile_size_cm:
        raise SystemExit(
            f"pitch-cm ({pitch_cm}) is smaller than tag-size-cm + 2*tag-white-border-cm "
            f"({tag_tile_size_cm}). Increase pitch-cm to keep white border around each tag."
        )

    inner_width_cm = args.width_cm - 2.0 * args.white_border_cm
    inner_height_cm = args.height_cm - 2.0 * args.white_border_cm

    if inner_width_cm + 1e-9 < tag_tile_size_cm or inner_height_cm + 1e-9 < tag_tile_size_cm:
        raise SystemExit(
            "After white border is reserved, inner area must still fit at least one full tag tile in both dimensions"
        )

    aruco_dict, max_tags = resolve_aruco_dict(args.family)

    width_px = cm_to_px(args.width_cm, args.dpi)
    height_px = cm_to_px(args.height_cm, args.dpi)
    tag_tile_px = cm_to_px(tag_tile_size_cm, args.dpi)
    marker_px = cm_to_px(args.tag_size_cm, args.dpi)

    if width_px < tag_tile_px or height_px < tag_tile_px:
        raise SystemExit("At current DPI, tag tile is larger than the canvas")
    if marker_px <= 0:
        raise SystemExit("At current DPI, tag-size-cm leaves no drawable marker pixels")

    canvas = np.full((height_px, width_px), 255, dtype=np.uint8)

    top_x_local = border_positions(inner_width_cm, tag_tile_size_cm, pitch_cm)
    side_y_local = side_positions_excluding_corners(inner_height_cm, tag_tile_size_cm, pitch_cm)

    placements_cm: List[Tuple[float, float]] = []

    # Clockwise: top -> right -> bottom -> left
    top_y = args.white_border_cm
    left_x = args.white_border_cm
    right_x = args.white_border_cm + inner_width_cm - tag_tile_size_cm
    bottom_y = args.white_border_cm + inner_height_cm - tag_tile_size_cm

    for x_local in top_x_local:
        placements_cm.append((args.white_border_cm + x_local, top_y))
    for y_local in side_y_local:
        placements_cm.append((right_x, args.white_border_cm + y_local))
    for x_local in reversed(top_x_local):
        placements_cm.append((args.white_border_cm + x_local, bottom_y))
    for y_local in reversed(side_y_local):
        placements_cm.append((left_x, args.white_border_cm + y_local))

    placements_csv: List[Tuple[int, float, float, int, int]] = []
    current_id = args.start_id

    for x_cm, y_cm in placements_cm:
        if current_id >= max_tags:
            if args.wrap_ids:
                tag_id = current_id % max_tags
            else:
                raise SystemExit(
                    f"Tag ID {current_id} exceeds family max ({max_tags - 1}). "
                    "Use --wrap-ids or lower --start-id."
                )
        else:
            tag_id = current_id

        x_px = cm_to_px(x_cm, args.dpi)
        y_px = cm_to_px(y_cm, args.dpi)

        # Clamp to keep full tag tile in canvas despite rounding.
        x_px = min(max(x_px, 0), width_px - tag_tile_px)
        y_px = min(max(y_px, 0), height_px - tag_tile_px)

        marker = make_marker(aruco_dict, tag_id, marker_px)

        marker_x_px = cm_to_px(x_cm + tag_white_border_cm, args.dpi)
        marker_y_px = cm_to_px(y_cm + tag_white_border_cm, args.dpi)

        # Keep marker fully inside its tag tile despite rounding.
        marker_x_px = min(max(marker_x_px, x_px), x_px + tag_tile_px - marker_px)
        marker_y_px = min(max(marker_y_px, y_px), y_px + tag_tile_px - marker_px)

        canvas[marker_y_px : marker_y_px + marker_px, marker_x_px : marker_x_px + marker_px] = marker
        placements_csv.append((tag_id, round(x_cm, 4), round(y_cm, 4), x_px, y_px))

        current_id += 1

    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.output.suffix.lower() == ".pdf":
        save_a4_pdf(
            args.output,
            canvas,
            args.width_cm,
            args.height_cm,
            args.dpi,
            landscape=args.a4_landscape,
        )
    else:
        save_printable_png(args.output, canvas, args.width_cm, args.height_cm, args.dpi)

    if args.layout_csv is not None:
        args.layout_csv.parent.mkdir(parents=True, exist_ok=True)
        write_layout_csv(args.layout_csv, placements_csv)

    print(f"Saved border image: {args.output}")
    print(
        f"family={args.family}, tags={len(placements_csv)}, "
        f"size={args.width_cm}cm x {args.height_cm}cm, tag_marker={args.tag_size_cm}cm, "
        f"pitch={round(pitch_cm, 6)}cm, white_border={args.white_border_cm}cm, "
        f"tag_white_ratio={round(tag_white_ratio, 6)}, "
        f"tag_white_border={round(tag_white_border_cm, 6)}cm, "
        f"tag_tile={round(tag_tile_size_cm, 6)}cm, output={args.output.suffix.lower()}"
    )
    if args.layout_csv is not None:
        print(f"Saved layout CSV: {args.layout_csv}")

    return 0


if __name__ == "__main__":
    sys.exit(main())
