"""
Long arm, single joint per shoulder (no welded dummy shoulders).

Loaded by: python generate_example_twolink_custom.py --config medium_single_redo
"""

import numpy as np
from pathlib import Path
import math

USE_WORLD_LINKS = True

# Marker orientation: use one of
#   "r": [nx,ny,nz]  quarter turns about world/body x,y,z (integers, legacy)
#   "quat": (w,x,y,z)  wxyz, MuJoCo
#   "euler_deg": [rx,ry,rz]  optional "euler_order" default "xyz" (MuJoCo mju_euler2Quat; lower/upper = intrinsic/fixed)
#   "euler" / "euler_rad"  same angles in radians
aruco_positions = {
    "base": {
        "pos": np.array([-32.86, -70, 14.2]) / 1000,
        "r": [0, 0, 0.0],
        "size": 50,  # Blender plane scale
        "size_is_blender_plane_scale": True,
        "aruco_id": 0, # todo fix this
        "aruco_grid_size": (3,3),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_base.png",
    },
    "shoulder": {
        # Given in base (parent link of shoulder) coordinates; generator converts to shoulder-local.
        "frame": "parent_link",
        "pos": np.array([93.9, -22.3, 44.02]) / 1000,
        "r": [1, 0, 0.0],  # 90,0,0
        "size": 15,  # Blender plane scale
        "size_is_blender_plane_scale": True,
        "aruco_id": 20,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_base_clamp.png",
    },
    "shoulder2": {
        # Given in base2 (parent link of shoulder2) coordinates; same size/grid as other clamps.
        "frame": "parent_link",
        "pos": np.array([94.4, 1.56, 64.3]) / 1000,
        "r": [0, 0, 0.0],  # 90,0,0
        "size": 15,  # same for future clamps
        "size_is_blender_plane_scale": True,
        "aruco_id": 24,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_link1_clamp.png",
    },
    "shoulder3": {
        # Given in base3 (parent link of shoulder3) coordinates; same size/grid as other clamps.
        "frame": "parent_link",
        "pos": np.array([103,.9,60]) / 1000,
        "r": [0, 0, 0.0],
        "size": 15,
        "size_is_blender_plane_scale": True,
        "aruco_id": 28,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_link2_clamp.png",
    },
    "shoulder4": {
        # Given in base4 (parent link of shoulder4) coordinates; same size/grid as other clamps.
        "frame": "parent_link",
        "pos": np.array([63.44, 0.92, 74.3]) / 1000,
        "r": [0, 0, 0.0],
        "size": 15,
        "size_is_blender_plane_scale": True,
        "aruco_id": 32,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_link3_clamp.png",
    },
    "shoulder5": {
        # Given in base5 (parent link of shoulder5) coordinates; same size/grid as other clamps.
        "frame": "parent_link",
        "pos": np.array([71.5, -21.4, 42.0]) / 1000,
        "euler_deg": [90, 0, 0],
        "size": 15,
        "size_is_blender_plane_scale": True,
        "aruco_id": 36,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_link4_clamp.png",
    },
    "shoulder6": {
        # Given in base6 (parent link of shoulder6) coordinates; same size/grid as other clamps.
        "frame": "parent_link",
        "pos": np.array([74.5, 18.0, 33.0]) / 1000,
        "euler_deg": [-90, 0, 270],
        "size": 15,
        "size_is_blender_plane_scale": True,
        "aruco_id": 40,
        "aruco_grid_size": (2, 2),
        "outer_margin_frac": 0.06,
        "texture_file": "aruco_link5_clamp.png",
    },
}


def _generate_aruco_texture(
    marker_id: int,
    texture_file: str,
    grid_size=(1, 1),
    outer_margin_frac: float = 0.0,
    px: int = 800,
    border_frac: float = 0.02,
) -> None:
    """Generate an ArUco marker/grid texture into ../assets for MuJoCo."""
    try:
        import cv2  # pyright: ignore[reportMissingImports]
    except Exception:
        return
    assets_dir = Path(__file__).resolve().parents[1] / "assets"
    assets_dir.mkdir(parents=True, exist_ok=True)
    out = assets_dir / texture_file
    dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250)
    gx, gy = int(grid_size[0]), int(grid_size[1])
    margin = max(0.0, min(0.45, float(outer_margin_frac)))
    inner_px = max(8, int(round(px * (1.0 - 2.0 * margin))))
    border_px = max(1, int(round(px * max(0.0, min(0.2, float(border_frac))))))

    if gx <= 1 and gy <= 1:
        marker = cv2.aruco.generateImageMarker(dictionary, int(marker_id), int(inner_px))
        canvas = np.full((int(px), int(px)), 255, dtype=np.uint8)
        off = (int(px) - int(inner_px)) // 2
        canvas[off : off + int(inner_px), off : off + int(inner_px)] = marker
        # Printable border around full marker+margin area.
        canvas[:border_px, :] = 0
        canvas[-border_px:, :] = 0
        canvas[:, :border_px] = 0
        canvas[:, -border_px:] = 0
        cv2.imwrite(str(out), canvas)
        return

    ids = np.arange(int(marker_id), int(marker_id) + gx * gy, dtype=np.int32)
    board = cv2.aruco.GridBoard((gx, gy), 0.8, 0.2, dictionary, ids)
    img = board.generateImage((int(inner_px), int(inner_px)), marginSize=0, borderBits=1)
    canvas = np.full((int(px), int(px)), 255, dtype=np.uint8)
    off = (int(px) - int(inner_px)) // 2
    canvas[off : off + int(inner_px), off : off + int(inner_px)] = img
    # Printable border around full marker+margin area.
    canvas[:border_px, :] = 0
    canvas[-border_px:, :] = 0
    canvas[:, :border_px] = 0
    canvas[:, -border_px:] = 0
    cv2.imwrite(str(out), canvas)


def _entry_board_size_m(entry: dict) -> float:
    """Match generator board-size semantics for printed physical size."""
    size_raw = float(entry.get("size", 50.0))
    board_size_m = size_raw if size_raw <= 2.0 else (size_raw / 1000.0)
    if bool(entry.get("size_is_blender_plane_scale", False)):
        board_size_m *= 2.0
    return board_size_m


def _generate_aruco_print_sheet(target_dpi: int = 300) -> None:
    """Create print sheets with true per-marker physical scale."""
    try:
        import cv2  # pyright: ignore[reportMissingImports]
    except Exception:
        return

    assets_dir = Path(__file__).resolve().parents[1] / "assets"
    assets_dir.mkdir(parents=True, exist_ok=True)

    entries = []
    for name, entry in aruco_positions.items():
        tex = entry.get("texture_file", f"aruco_{name}.png")
        p = assets_dir / tex
        img = cv2.imread(str(p), cv2.IMREAD_GRAYSCALE)
        if img is None:
            continue
        entries.append((name, entry, img))

    if not entries:
        return

    px_per_m = float(target_dpi) / 0.0254
    marker_px = max(max(8, int(round(_entry_board_size_m(entry) * px_per_m))) for _, entry, _ in entries)
    title_h = 56
    label_h = 116
    pad = 24
    tile_w = marker_px + 2 * pad
    tile_h = title_h + marker_px + label_h + 2 * pad
    cols = min(3, max(1, len(entries)))
    rows = int(math.ceil(len(entries) / cols))
    sheet = np.full((rows * tile_h + pad, cols * tile_w + pad), 255, dtype=np.uint8)

    for idx, (name, entry, img) in enumerate(entries):
        r = idx // cols
        c = idx % cols
        x0 = c * tile_w + pad
        y0 = r * tile_h + pad

        # Header and marker.
        cv2.putText(
            sheet,
            f"{name}",
            (x0 + 8, y0 + 34),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.9,
            0,
            2,
            cv2.LINE_AA,
        )
        target_marker_px = max(8, int(round(_entry_board_size_m(entry) * px_per_m)))
        if img.shape[0] != target_marker_px:
            img_draw = cv2.resize(img, (target_marker_px, target_marker_px), interpolation=cv2.INTER_NEAREST)
        else:
            img_draw = img
        off_x = x0 + (tile_w - img_draw.shape[1]) // 2
        off_y = y0 + title_h
        sheet[off_y : off_y + img_draw.shape[0], off_x : off_x + img_draw.shape[1]] = img_draw

        # Annotation lines for print/use.
        gid = int(entry.get("aruco_id", 0))
        gx, gy = entry.get("aruco_grid_size", (1, 1))
        board_mm = _entry_board_size_m(entry) * 1000.0
        cv2.putText(
            sheet,
            f"ids: {gid}-{gid + int(gx) * int(gy) - 1}  grid: {gx}x{gy}",
            (x0 + 8, off_y + img_draw.shape[0] + 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.65,
            0,
            2,
            cv2.LINE_AA,
        )
        cv2.putText(
            sheet,
            f"board: {board_mm:.1f} mm  @ {target_dpi} DPI",
            (x0 + 8, off_y + img_draw.shape[0] + 62),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            0,
            2,
            cv2.LINE_AA,
        )
        cv2.putText(
            sheet,
            f"texture: {entry.get('texture_file', f'aruco_{name}.png')}",
            (x0 + 8, off_y + img_draw.shape[0] + 94),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            0,
            1,
            cv2.LINE_AA,
        )

    png_out = assets_dir / "aruco_print_sheet_medium_single_redo.png"
    cv2.imwrite(str(png_out), sheet)

    try:
        from PIL import Image  # pyright: ignore[reportMissingImports]

        # Save overview PNG with embedded DPI metadata.
        Image.fromarray(sheet).save(str(png_out), dpi=(float(target_dpi), float(target_dpi)))

        # Build true-scale A4 multi-page printable PDF to avoid printer "fit to page" shrinking.
        a4_w_px = int(round((210.0 / 25.4) * target_dpi))
        a4_h_px = int(round((297.0 / 25.4) * target_dpi))
        margin_px = int(round(10.0 / 25.4 * target_dpi))  # 10mm margins
        gap_px = int(round(6.0 / 25.4 * target_dpi))      # 6mm gaps
        title_h_px = int(round(10.0 / 25.4 * target_dpi))
        line_h_px = int(round(6.0 / 25.4 * target_dpi))
        foot_h_px = int(round(4.0 / 25.4 * target_dpi))
        label_block_px = title_h_px + 3 * line_h_px + foot_h_px

        pages = []
        page = np.full((a4_h_px, a4_w_px), 255, dtype=np.uint8)
        x = margin_px
        y = margin_px
        row_h = 0

        def _new_page():
            return np.full((a4_h_px, a4_w_px), 255, dtype=np.uint8)

        calib_mm = 50.0
        calib_px = int(round(calib_mm / 25.4 * target_dpi))
        calib_pad = int(round(3.0 / 25.4 * target_dpi))

        def _draw_calibration(g):
            # 50mm scale check to detect hidden print scaling.
            x0 = a4_w_px - margin_px - calib_px
            y0 = margin_px
            x1 = x0 + calib_px
            y1 = y0 + calib_px
            cv2.rectangle(g, (x0, y0), (x1, y1), 0, 2)
            cv2.line(g, (x0, y1 + calib_pad), (x1, y1 + calib_pad), 0, 2)
            cv2.putText(
                g,
                "50 mm calibration (must measure exactly 50mm)",
                (x0 - int(round(30.0 / 25.4 * target_dpi)), y1 + calib_pad + int(round(5.0 / 25.4 * target_dpi))),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.42,
                0,
                1,
                cv2.LINE_AA,
            )

        _draw_calibration(page)

        for name, entry, img in entries:
            board_px = max(8, int(round(_entry_board_size_m(entry) * px_per_m)))
            img_draw = cv2.resize(img, (board_px, board_px), interpolation=cv2.INTER_NEAREST) if img.shape[0] != board_px else img
            tile_w = img_draw.shape[1]
            tile_h = label_block_px + img_draw.shape[0]

            # New row if needed.
            if x + tile_w > (a4_w_px - margin_px):
                x = margin_px
                y += row_h + gap_px
                row_h = 0
            # New page if needed.
            if y + tile_h > (a4_h_px - margin_px):
                pages.append(page)
                page = _new_page()
                _draw_calibration(page)
                x = margin_px
                y = margin_px
                row_h = 0

            gid = int(entry.get("aruco_id", 0))
            gx, gy = entry.get("aruco_grid_size", (1, 1))
            board_mm = _entry_board_size_m(entry) * 1000.0
            cv2.putText(page, f"{name}", (x, y + int(0.8 * title_h_px)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, 0, 2, cv2.LINE_AA)
            cv2.putText(page, f"ids: {gid}-{gid + int(gx) * int(gy) - 1}  grid: {gx}x{gy}", (x, y + title_h_px + int(0.75 * line_h_px)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 0, 1, cv2.LINE_AA)
            cv2.putText(page, f"board: {board_mm:.1f} mm @ {target_dpi} DPI", (x, y + title_h_px + int(1.8 * line_h_px)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, 0, 1, cv2.LINE_AA)
            cv2.putText(page, f"texture: {entry.get('texture_file', f'aruco_{name}.png')}", (x, y + title_h_px + int(2.85 * line_h_px)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, 0, 1, cv2.LINE_AA)

            iy = y + label_block_px
            page[iy : iy + img_draw.shape[0], x : x + img_draw.shape[1]] = img_draw

            x += tile_w + gap_px
            row_h = max(row_h, tile_h)

        pages.append(page)
        pil_pages = [Image.fromarray(p).convert("RGB") for p in pages]
        pil_pages[0].save(
            str(assets_dir / "aruco_print_sheet_medium_single_redo.pdf"),
            "PDF",
            resolution=float(target_dpi),
            save_all=True,
            append_images=pil_pages[1:],
        )
    except Exception:
        pass


for _name, _entry in aruco_positions.items():
    _generate_aruco_texture(
        _entry.get("aruco_id", 0),
        _entry.get("texture_file", f"aruco_{_name}.png"),
        _entry.get("aruco_grid_size", (1, 1)),
        _entry.get("outer_margin_frac", 0.0),
    )

_generate_aruco_print_sheet()

WORLD_LINKS = [
    {"pos": np.array([0, 0, 0]) / 1000, "r": [0, 0, 0.0]},
    {"pos": np.array([100,42,78.46]) / 1000, "r": [1, -1, 0]},
    {"pos": np.array([117,40,240]) / 1000, "r": [1, 0, 0]},
    {"pos": np.array([260,40,240,]) / 1000, "r": [1, 0, 0]},
    {"pos": np.array([325,0,199]) / 1000, "r": [0, 0, 0]},
    {"pos": np.array([443,-5.5,211]) / 1000, "r": [1, 1, -1]},
    {"pos": np.array([443,32.5,165.5]) / 1000, "r": [1, 1, -1]},
    #{"pos": np.array([418.6,41.9,266]) / 1000, "r": [1, 1, 0]},
    #{"pos": np.array([466.7,3,197]) / 1000, "r": [1, 1, -1]},
    #{"pos": np.array([466.7,41,151.56]) / 1000, "r": [1, 1, -1]},
]

# Per link index (0=base, 1=base2, …). None = omit holder/clamp geoms for that link.
HOLDER_MESH_FILES = [
    "arm_redo_thin_long_pretty_medium_size_base.stl",
    None,#"arm_redo_thin_long_pretty_medium_size_base.stl",
    None,
    None,
    None,
    None
]

# Clamp STLs in world frame at home: p_w = T_wl @ T_ls(q_i) @ T_clamp @ verts; geom = inv(T_ls(qi)) @ inv(T_wl).
USE_LINK_CLAMP_COORD_FRAMES = False
# Holder STLs in world frame: p_w = T_wl @ T_holder @ verts; geom = inv(T_holder) @ inv(T_wl).
USE_LINK_HOLDER_COORD_FRAMES = False

CLAMP_MESH_FILES = [
    "arm_redo_thin_long_pretty_medium_size_base_clamp.stl",
    "link2_)stronger.stl",
    "arm_redo_thin_long_pretty_medium_size_link2_clamp.stl",
    "arm_redo_thin_long_pretty_medium_size_link3_clamp.stl",
    "arm_redo_thin_long_pretty_medium_size_link4_clamp_redo.stl",
    #"arm_redo_thin_long_pretty_medium_size_link5_clamp.stl",
    "eef_base_with_handle.stl",
    "eef_thumb_with_handle.stl",
]

# Welded shoulders (no hinge): set of shoulder indices 0=shoulder, 1=shoulder2, …
SHOULDER_INDICES_WITHOUT_JOINT = set()

# Used when USE_WORLD_LINKS is False
LINK_POSITIONS = [
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
]

LINK_ROTATIONS = [
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
    [0, 0, 0],
]

DEFAULT_JOINT_QPOS = [0.0, -1.5, 1.5, 0.75, 0.0, 0.0]

EQUAL_JOINT_PAIRS = []

LINK_MATERIALS = [
    "yellow_mat",
    "red_mat",
    "green_mat",
    "blue_mat",
    "orange_mat",
]