{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "NDArray = np.ndarray\n",
    "\n",
    "# Order of detection result\n",
    "# 0,   1,     2,    3\n",
    "# TL,  TR,    BR,   BL\n",
    "# RED, GREEN, BLUE, YELLOW\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DiamondBoardParameter:\n",
    "    marker_leghth: float\n",
    "    \"\"\"\n",
    "    the ArUco marker length in meter\n",
    "    \"\"\"\n",
    "    chess_length: float\n",
    "    \"\"\"\n",
    "    the length of the chess board in meter\n",
    "    \"\"\"\n",
    "    border_length: float = 0.01\n",
    "    \"\"\"\n",
    "    border_length in m, default is 1cm\n",
    "    \"\"\"\n",
    "\n",
    "    @property\n",
    "    def marker_border_length(self):\n",
    "        assert self.chess_length > self.marker_leghth\n",
    "        return (self.chess_length - self.marker_leghth) / 2\n",
    "\n",
    "    @property\n",
    "    def total_side_length(self):\n",
    "        assert self.chess_length > self.marker_leghth\n",
    "        return self.marker_border_length * 2 + self.chess_length * 3\n",
    "\n",
    "\n",
    "# 9mm + 127mm + 127mm (97mm marker) + 127mm + 10mm\n",
    "# i.e. marker boarder = 127mm - 97mm = 30mm (15mm each side)\n",
    "Point2D = tuple[float, float]\n",
    "Quad2D = tuple[Point2D, Point2D, Point2D, Point2D]\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ArUcoMarker2D:\n",
    "    id: int\n",
    "    corners: Quad2D\n",
    "    params: DiamondBoardParameter\n",
    "\n",
    "    @property\n",
    "    def np_corners(self):\n",
    "        \"\"\"\n",
    "        returns corners in numpy array\n",
    "        (4, 2) shape\n",
    "        \"\"\"\n",
    "        return np.array(self.corners, dtype=np.float32)\n",
    "\n",
    "\n",
    "# let's let TL be the origin\n",
    "def generate_diamond_corners(\n",
    "    ids: tuple[int, int, int, int], params: DiamondBoardParameter\n",
    "):\n",
    "    \"\"\"\n",
    "    A diamond chess board, which could be count as a kind of ChArUco board\n",
    "\n",
    "    C | 0 | C\n",
    "    ---------\n",
    "    1 | C | 2\n",
    "    ---------\n",
    "    C | 3 | C\n",
    "\n",
    "    where C is the chess box, and 0, 1, 2, 3 are the markers (whose ids are passed in order)\n",
    "\n",
    "    Args:\n",
    "        ids: a tuple of 4 ids of the markers\n",
    "        params: DiamondBoardParameter\n",
    "    \"\"\"\n",
    "\n",
    "    def tl_to_square(tl_x: float, tl_y: float, side_length: float) -> Quad2D:\n",
    "        return (\n",
    "            (tl_x, tl_y),\n",
    "            (tl_x + side_length, tl_y),\n",
    "            (tl_x + side_length, tl_y + side_length),\n",
    "            (tl_x, tl_y + side_length),\n",
    "        )\n",
    "\n",
    "    tl_0_x = params.border_length + params.chess_length + params.marker_border_length\n",
    "    tl_0_y = params.border_length + params.marker_border_length\n",
    "\n",
    "    tl_1_x = params.border_length + params.marker_border_length\n",
    "    tl_1_y = params.border_length + params.chess_length + params.marker_border_length\n",
    "\n",
    "    tl_2_x = (\n",
    "        params.border_length + params.chess_length * 2 + params.marker_border_length\n",
    "    )\n",
    "    tl_2_y = tl_1_y\n",
    "\n",
    "    tl_3_x = params.border_length + params.chess_length + params.marker_border_length\n",
    "    tl_3_y = (\n",
    "        params.border_length + params.chess_length * 2 + params.marker_border_length\n",
    "    )\n",
    "    return (\n",
    "        ArUcoMarker2D(ids[0], tl_to_square(tl_0_x, tl_0_y, params.marker_leghth), params),\n",
    "        ArUcoMarker2D(ids[1], tl_to_square(tl_1_x, tl_1_y, params.marker_leghth), params),\n",
    "        ArUcoMarker2D(ids[2], tl_to_square(tl_2_x, tl_2_y, params.marker_leghth), params),\n",
    "        ArUcoMarker2D(ids[3], tl_to_square(tl_3_x, tl_3_y, params.marker_leghth), params),\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = DiamondBoardParameter(0.097, 0.127)\n",
    "markers = generate_diamond_corners((16, 17, 18, 19), params)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.gca()\n",
    "ax.set_xlim((0, params.total_side_length))\n",
    "ax.set_ylim((0, params.total_side_length)) # type: ignore\n",
    "ax.set_aspect(\"equal\")\n",
    "# set origin to top-left (from bottom-left)\n",
    "ax.invert_yaxis()\n",
    "ax.xaxis.set_ticks_position('top')\n",
    "\n",
    "for marker in markers:\n",
    "    plt.plot(*marker.np_corners.T, \"o-\", label=str(marker.id))\n",
    "    for i, (x, y) in enumerate(marker.corners):\n",
    "        ax.text(x, y, str(i))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence\n",
    "import plotly.graph_objects as go\n",
    "import awkward as ak\n",
    "import cv2\n",
    "from cv2.typing import MatLike\n",
    "\n",
    "\n",
    "def transform_point(matrix: MatLike, point: MatLike):\n",
    "    assert matrix.shape == (4, 4)\n",
    "    assert point.shape == (3,)\n",
    "\n",
    "    # Lift point to 4D\n",
    "    homogeneous_point = np.array([point[0], point[1], point[2], 1])\n",
    "    # Apply transformation\n",
    "    transformed = matrix @ homogeneous_point\n",
    "    # Project back to 3D if w != 1\n",
    "    if transformed[3] != 1:\n",
    "        transformed = transformed / transformed[3]\n",
    "    return transformed[:3]\n",
    "\n",
    "\n",
    "class DiamondPlane3D:\n",
    "    _ids: NDArray\n",
    "    \"\"\"\n",
    "    (n,)\n",
    "    \"\"\"\n",
    "    _corners: NDArray\n",
    "    \"\"\"\n",
    "    (n, 4, 3)\n",
    "    \"\"\"\n",
    "    _transform_matrix: NDArray\n",
    "    \"\"\"\n",
    "    4x4 transformation matrix\n",
    "    \"\"\"\n",
    "    _normal_vector: NDArray\n",
    "    \"\"\"\n",
    "    (2, 3)\n",
    "    start (the center of the plane) and end (the normal vector), length 1\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, items: Sequence[ArUcoMarker2D]):\n",
    "        self._ids = np.array([item.id for item in items])\n",
    "        # (n, 4, 2)\n",
    "        corners_2d = np.array([item.np_corners for item in items])\n",
    "        # (n, 4, 3)\n",
    "        self._corners = np.concatenate(\n",
    "            [corners_2d, np.zeros((corners_2d.shape[0], 4, 1))], axis=-1\n",
    "        )\n",
    "        self._transform_matrix = np.eye(4)\n",
    "\n",
    "        def center(items: Sequence[ArUcoMarker2D]):\n",
    "            return np.mean([item.np_corners for item in items], axis=(0, 1))\n",
    "\n",
    "        c = center(items)\n",
    "        assert c.shape == (2,)\n",
    "        self._normal_vector = np.array([(c[0], c[1], 0), (c[0], c[1], 0.1)])\n",
    "\n",
    "    @property\n",
    "    def ids(self):\n",
    "        return self._ids\n",
    "\n",
    "    @property\n",
    "    def corners(self):\n",
    "        return self._corners\n",
    "\n",
    "    @property\n",
    "    def transform_matrix(self):\n",
    "        return self._transform_matrix\n",
    "\n",
    "    @property\n",
    "    def transformed_corners(self):\n",
    "        def g():\n",
    "            for corner in self.corners:\n",
    "                yield np.array(\n",
    "                    [transform_point(self.transform_matrix, c) for c in corner]\n",
    "                )\n",
    "\n",
    "        return np.array(list(g()))\n",
    "\n",
    "    @property\n",
    "    def transformed_normal_vector(self):\n",
    "        def g():\n",
    "            for v in self._normal_vector:\n",
    "                yield transform_point(self.transform_matrix, v)\n",
    "\n",
    "        return np.array(list(g()))\n",
    "\n",
    "    @property\n",
    "    def transformed_geometry_center(self):\n",
    "        return np.mean(self.transformed_corners, axis=(0, 1))\n",
    "\n",
    "    def local_rotate(self, angle: float, axis: NDArray):\n",
    "        \"\"\"\n",
    "        rotate the plane by angle (in radian) around local center\n",
    "\n",
    "        Args:\n",
    "            angle: in radian\n",
    "            axis: (3,)\n",
    "\n",
    "        change basis to local basis, rotate, then change back\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def rotate(self, angle: float, axis: NDArray):\n",
    "        \"\"\"\n",
    "        rotate the plane by angle (in radian) around the axis\n",
    "        \"\"\"\n",
    "        assert axis.shape == (3,)\n",
    "        rot_mat = cv2.Rodrigues(axis * angle)[0]\n",
    "        self._transform_matrix[:3, :3] = np.dot(rot_mat, self._transform_matrix[:3, :3])\n",
    "\n",
    "    def translate(self, vec: NDArray):\n",
    "        \"\"\"\n",
    "        translate the plane by vec\n",
    "        \"\"\"\n",
    "        assert vec.shape == (3,)\n",
    "        self._transform_matrix[:3, 3] += vec\n",
    "\n",
    "    def set_transform_matrix(self, mat: NDArray):\n",
    "        assert mat.shape == (4, 4)\n",
    "        self._transform_matrix = mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "plane_a = DiamondPlane3D(markers)\n",
    "\n",
    "OFFSET = 0.000\n",
    "markers_b = generate_diamond_corners((20, 21, 22, 23), params)\n",
    "plane_b = DiamondPlane3D(markers_b)\n",
    "# plane_b.translate(np.array([0, 0, 0.1]))\n",
    "plane_b.rotate(np.pi/2, np.array([1, 0, 0]))\n",
    "plane_b.rotate(np.pi, np.array([0, 0, 1]))\n",
    "tmp_c = plane_b.transformed_geometry_center\n",
    "plane_b.translate(-tmp_c)\n",
    "plane_b.rotate(np.pi, np.array([0, 1, 0]))\n",
    "plane_b.translate(tmp_c)\n",
    "plane_b.translate(np.array([0, 0, params.total_side_length]))\n",
    "plane_b.translate(np.array([0, 0, -OFFSET]))\n",
    "# OFFSET for plane_b\n",
    "# plane_b.translate(np.array([0, 0.001, 0]))\n",
    "\n",
    "markers_c = generate_diamond_corners((24, 25, 26, 27), params)\n",
    "plane_c = DiamondPlane3D(markers_c)\n",
    "tmp = plane_c.transformed_geometry_center\n",
    "plane_c.translate(-tmp)\n",
    "plane_c.rotate(-np.pi/2, np.array([0, 0, 1]))\n",
    "plane_c.translate(tmp)\n",
    "plane_c.translate(np.array([0, params.total_side_length-params.border_length, 0]))\n",
    "plane_c.rotate(np.pi/2, np.array([0, 1, 0]))\n",
    "plane_c.translate(np.array([0, 0, params.total_side_length]))\n",
    "plane_c.translate(np.array([0, 0, -OFFSET]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "t_corners_a = plane_a.transformed_corners\n",
    "for i, corners in enumerate(t_corners_a):\n",
    "    fig.add_trace(\n",
    "        go.Scatter3d(\n",
    "            x=corners[:, 0],\n",
    "            y=corners[:, 1],\n",
    "            z=corners[:, 2],\n",
    "            mode=\"markers+lines+text\",\n",
    "            text=list(map(lambda x: f\"{plane_a.ids[i]}:{x}\", range(4))),\n",
    "            textposition=\"middle center\",\n",
    "            name=str(plane_a.ids[i]),\n",
    "            marker=dict(size=1),\n",
    "        )\n",
    "    )\n",
    "\n",
    "# normal vector\n",
    "fig.add_trace(\n",
    "    go.Scatter3d(\n",
    "        x=plane_a.transformed_normal_vector[:, 0],\n",
    "        y=plane_a.transformed_normal_vector[:, 1],\n",
    "        z=plane_a.transformed_normal_vector[:, 2],\n",
    "        mode=\"markers+lines\",\n",
    "        name=\"normal_a\",\n",
    "        marker=dict(size=2),\n",
    "    )\n",
    ")\n",
    "\n",
    "t_corners_b = plane_b.transformed_corners\n",
    "for i, corners in enumerate(t_corners_b):\n",
    "    fig.add_trace(\n",
    "        go.Scatter3d(\n",
    "            x=corners[:, 0],\n",
    "            y=corners[:, 1],\n",
    "            z=corners[:, 2],\n",
    "            mode=\"markers+lines+text\",\n",
    "            text=list(map(lambda x: f\"{plane_b.ids[i]}:{x}\", range(4))),\n",
    "            textposition=\"middle center\",\n",
    "            name=str(plane_b.ids[i]),\n",
    "            marker=dict(size=1),\n",
    "        )\n",
    "    )\n",
    "fig.add_trace(\n",
    "    go.Scatter3d(\n",
    "        x=plane_b.transformed_normal_vector[:, 0],\n",
    "        y=plane_b.transformed_normal_vector[:, 1],\n",
    "        z=plane_b.transformed_normal_vector[:, 2],\n",
    "        mode=\"markers+lines\",\n",
    "        name=\"normal_b\",\n",
    "        marker=dict(size=2),\n",
    "    )\n",
    ")\n",
    "\n",
    "t_corners_c = plane_c.transformed_corners\n",
    "for i, corners in enumerate(t_corners_c):\n",
    "    fig.add_trace(\n",
    "        go.Scatter3d(\n",
    "            x=corners[:, 0],\n",
    "            y=corners[:, 1],\n",
    "            z=corners[:, 2],\n",
    "            mode=\"markers+lines+text\",\n",
    "            text=list(map(lambda x: f\"{plane_c.ids[i]}:{x}\", range(4))),\n",
    "            name=str(plane_c.ids[i]),\n",
    "            marker=dict(size=1),\n",
    "        )\n",
    "    )\n",
    "fig.add_trace(\n",
    "    go.Scatter3d(\n",
    "        x=plane_c.transformed_normal_vector[:, 0],\n",
    "        y=plane_c.transformed_normal_vector[:, 1],\n",
    "        z=plane_c.transformed_normal_vector[:, 2],\n",
    "        mode=\"markers+lines\",\n",
    "        textposition=\"middle center\",\n",
    "        name=\"normal_c\",\n",
    "        marker=dict(size=2),\n",
    "    )\n",
    ")\n",
    "\n",
    "# fig.update_layout(\n",
    "#     scene=dict(\n",
    "#         aspectmode=\"cube\",\n",
    "#         yaxis_autorange=\"reversed\",\n",
    "#     )\n",
    "# )\n",
    "\n",
    "fig.update_layout(\n",
    "    scene=dict(\n",
    "        aspectmode='cube',\n",
    "        xaxis=dict(range=[-0.1, params.total_side_length]),\n",
    "        yaxis=dict(range=[params.total_side_length, -0.1]),\n",
    "        zaxis=dict(range=[-0.1, params.total_side_length]),\n",
    "    )\n",
    ")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import awkward as ak\n",
    "from awkward import Record as AwkwardRecord, Array as AwkwardArray\n",
    "\n",
    "coords = AwkwardArray(\n",
    "    [\n",
    "        {\n",
    "            \"name\": \"a\",\n",
    "            \"ids\": plane_a.ids,\n",
    "            \"corners\": t_corners_a,\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"b\",\n",
    "            \"ids\": plane_b.ids,\n",
    "            \"corners\": t_corners_b,\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"c\",\n",
    "            \"ids\": plane_c.ids,\n",
    "            \"corners\": t_corners_c,\n",
    "        },\n",
    "    ]\n",
    ")\n",
    "display(coords)\n",
    "_ = ak.to_parquet(coords, \"output/object_points.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import cast\n",
    "total_ids = cast(NDArray, ak.to_numpy(coords[\"ids\"])).flatten()\n",
    "total_corners = cast(NDArray, ak.to_numpy(coords[\"corners\"])).reshape(-1, 4, 3)\n",
    "#display(total_ids, total_corners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dict(zip(total_ids, total_corners))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_ids = np.concatenate([plane_a.ids, plane_b.ids, plane_c.ids])\n",
    "total_corners = np.concatenate([t_corners_a, t_corners_b, t_corners_c])\n",
    "id_corner_map: dict[int, NDArray] = dict(zip(total_ids, total_corners))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}