"""Generate the large calib board and hook it into MuJoCo.

Outputs (into ./calib_board_out/ + the arm model dir):
  calib_board_grid.png        exact 8px/mm board texture (true 35mm markers)
  calib_board_print_page.png  printable page + a 50mm scale bar to verify print scale
  calib_board.xml             standalone board plane + floor + light (for localization)
  <armdir>/arm_calib_board.xml   arm_umi_v2 + the calib board plane at BOARD_POSE (AR overlay)

Run where aruco_plane.obj and the arm model live (cad tree / mj_arm_v2).
"""
import os
import shutil
import sys
import xml.etree.ElementTree as ET

import numpy as np
import cv2

from aruco_boards import DICT, board as mkboard
from calib_board import CALIB_BOARD, CONTENT_MM, CANVAS_MM, PX_PER_MM, MARKER_MM, BOARD_POSE_XYZ

ARM_XML = sys.argv[1] if len(sys.argv) > 1 else "mj_arm_v2/arm_umi_v2.xml"
PLANE_OBJ = sys.argv[2] if len(sys.argv) > 2 else "mj_arm_v2/aruco_plane.obj"
OUT = "calib_board_out"
os.makedirs(OUT, exist_ok=True)
ppm = PX_PER_MM

# --- board texture at exact scale ---
cw, ch = int(round(CANVAS_MM[0] * ppm)), int(round(CANVAS_MM[1] * ppm))
margin_px = int(round(CALIB_BOARD["margin_m"] * 1000 * ppm))
gb = mkboard(CALIB_BOARD)
grid = gb.generateImage((cw, ch), marginSize=margin_px, borderBits=1)
grid_path = os.path.join(OUT, "calib_board_grid.png")
cv2.imwrite(grid_path, grid)

det = cv2.aruco.ArucoDetector(DICT, cv2.aruco.DetectorParameters())
corners, ids, _ = det.detectMarkers(cv2.cvtColor(grid, cv2.COLOR_GRAY2BGR) if grid.ndim == 2 else grid)
sides = [np.linalg.norm(c[0][0] - c[0][1]) for c in corners]
print(f"grid {cw}x{ch}px  {0 if ids is None else len(ids)} markers ids "
      f"{sorted(int(i) for i in ids.flatten()) if ids is not None else []}")
print(f"  median marker side {np.median(sides):.1f}px  (expect {MARKER_MM*ppm:.0f}px = {MARKER_MM}mm) "
      f"-> {np.median(sides)/ppm:.2f}mm")

# --- print page: grid + label + 50mm scale bar ---
page = cv2.cvtColor(grid, cv2.COLOR_GRAY2BGR)
strip = np.full((90, cw, 3), 255, np.uint8)
cv2.putText(strip, f"CALIB BOARD  {int(CANVAS_MM[0])}x{int(CANVAS_MM[1])}mm  {int(MARKER_MM)}mm markers  "
            f"ids {CALIB_BOARD['ids'][0]}-{CALIB_BOARD['ids'][-1]}  (print at 100%)",
            (12, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
bar = int(round(50 * ppm))
cv2.line(strip, (12, 62), (12 + bar, 62), (0, 0, 0), 3)
cv2.putText(strip, "50mm (measure to verify scale)", (12 + bar + 10, 68), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
page = np.vstack([page, strip])
cv2.rectangle(page, (1, 1), (cw - 2, page.shape[0] - 2), (0, 0, 0), 1)
cv2.imwrite(os.path.join(OUT, "calib_board_print_page.png"), page)

# --- embed real DPI + make an A4 PDF that prints at true size ---
# cv2 writes no DPI tag, so viewers assume 72dpi and "100%" balloons the image. Tag it at
# the true 8px/mm = 203.2dpi, and lay the page onto an A4 sheet as a PDF (foolproof print).
from PIL import Image  # noqa: E402
DPI = PX_PER_MM * 25.4
for f in ("calib_board_grid.png", "calib_board_print_page.png"):
    Image.open(os.path.join(OUT, f)).save(os.path.join(OUT, f), dpi=(DPI, DPI))
a4 = Image.new("RGB", (int(round(210 * PX_PER_MM)), int(round(297 * PX_PER_MM))), "white")
pp = Image.open(os.path.join(OUT, "calib_board_print_page.png")).convert("RGB")
a4.paste(pp, ((a4.width - pp.width) // 2, (a4.height - pp.height) // 2))
a4.save(os.path.join(OUT, "calib_board_a4.png"), dpi=(DPI, DPI))
a4.save(os.path.join(OUT, "calib_board.pdf"), "PDF", resolution=DPI)
print(f"tagged PNGs at {DPI:.1f} dpi; wrote calib_board.pdf (A4, prints at true size)")

# --- standalone MuJoCo XML ---
sx, sy = CANVAS_MM[0] / 300.0, CANVAS_MM[1] / 300.0   # aruco_plane.obj base is 0.3m
shutil.copy(PLANE_OBJ, os.path.join(OUT, "aruco_plane.obj"))
open(os.path.join(OUT, "calib_board.xml"), "w").write(f'''<mujoco model="calib_board">
  <compiler autolimits="true"/>
  <visual><global offwidth="1600" offheight="1600"/></visual>
  <asset>
    <texture name="calib" type="2d" file="calib_board_grid.png"/>
    <material name="calib" texture="calib"/>
    <mesh name="calib_plane" file="aruco_plane.obj" scale="{sx:.5f} {sy:.5f} 0.001" inertia="shell"/>
    <texture name="grid" type="2d" builtin="checker" rgb1="0.22 0.27 0.32" rgb2="0.16 0.2 0.24" width="300" height="300"/>
    <material name="floor" texture="grid" texuniform="true" texrepeat="4 4"/>
  </asset>
  <worldbody>
    <light pos="0 0 1.2" dir="0 0 -1" diffuse="0.9 0.9 0.9"/>
    <geom type="plane" size="1 1 0.05" pos="0 0 0" material="floor"/>
    <geom type="mesh" mesh="calib_plane" material="calib" pos="0 0 0.001" contype="0" conaffinity="0"/>
  </worldbody>
</mujoco>
''')

# --- composite arm + calib board ---
tree = ET.parse(ARM_XML)
root = tree.getroot()
asset = root.find("asset")
ET.SubElement(asset, "texture", {"name": "calib", "type": "2d", "file": "calib_board_grid.png"})
ET.SubElement(asset, "material", {"name": "calib", "texture": "calib"})
ET.SubElement(asset, "mesh", {"name": "calib_plane", "file": "aruco_plane.obj",
                              "scale": f"{sx:.5f} {sy:.5f} 0.001", "inertia": "shell"})
bx, by, bz = BOARD_POSE_XYZ
ET.SubElement(root.find("worldbody"), "geom",
              {"type": "mesh", "mesh": "calib_plane", "material": "calib",
               "pos": f"{bx} {by} {bz}", "contype": "0", "conaffinity": "0"})
ET.indent(tree, "  ")
arm_out = os.path.join(os.path.dirname(ARM_XML), "arm_calib_board.xml")
tree.write(arm_out, encoding="unicode")
shutil.copy(grid_path, os.path.join(os.path.dirname(ARM_XML), "calib_board_grid.png"))
print(f"wrote {OUT}/ (grid, print page, calib_board.xml) + {arm_out}")