{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fefb950-9158-4c62-b593-cda353ff5db1",
   "metadata": {},
   "source": [
    "# Demo of `LaPDXYTransform`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bef64d2-1541-4dec-ac10-ebcf4cffe4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c23fe0-5407-40b9-a998-6f1581d6eb6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import sys\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = [10.5, 0.56 * 10.5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e25f18e-6ce0-48b2-82ac-27c69b006a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from bapsf_motion.transform import LaPDXYTransform\n",
    "except ModuleNotFoundError:\n",
    "    from pathlib import Path\n",
    "\n",
    "    HERE = Path().cwd()\n",
    "    BAPSF_MOTION = (HERE / \"..\" / \"..\" / \"..\" ).resolve()\n",
    "    sys.path.append(str(BAPSF_MOTION))\n",
    "    \n",
    "    from bapsf_motion.transform import LaPDXYTransform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "120c8907-672f-4915-aa03-506dd91b1d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = LaPDXYTransform(\n",
    "    (\"x\", \"y\"),\n",
    "    pivot_to_center=57.288,\n",
    "    pivot_to_drive=134.0,\n",
    "    pivot_to_feedthru=21.6,\n",
    "    # probe_axis_offset=10.00125,\n",
    "    probe_axis_offset=20.16125,\n",
    "    droop_correct=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e61105e-3788-4edb-a54c-9a582f04f745",
   "metadata": {},
   "outputs": [],
   "source": [
    "figwidth, figheight = plt.rcParams[\"figure.figsize\"]\n",
    "figwidth = 1.4 * figwidth\n",
    "figheight = 2.0 * figheight\n",
    "fig, axs = plt.subplots(2, 3, figsize=[figwidth, figheight])\n",
    "\n",
    "axs[0,0].set_xlabel(\"MSpace X\")\n",
    "axs[0,0].set_ylabel(\"MSpace Y\")\n",
    "axs[0,1].set_xlabel(\"Drive X\")\n",
    "axs[0,1].set_ylabel(\"Drive Y\")\n",
    "axs[0,2].set_xlabel(\"MSpace X\")\n",
    "axs[0,2].set_ylabel(\"MSpace Y\")\n",
    "\n",
    "points = np.zeros((40, 2))\n",
    "points[0:10, 0] = np.linspace(-5, 5, num=10, endpoint=False)\n",
    "points[0:10, 1] = 5 * np.ones(10)\n",
    "points[10:20, 0] = 5 * np.ones(10)\n",
    "points[10:20, 1] = np.linspace(5, -5, num=10, endpoint=False)\n",
    "points[20:30, 0] = np.linspace(5, -5, num=10, endpoint=False)\n",
    "points[20:30, 1] = -5 * np.ones(10)\n",
    "points[30:40, 0] = -5 * np.ones(10)\n",
    "points[30:40, 1] = np.linspace(-5, 5, num=10, endpoint=False)\n",
    "\n",
    "dpoints = tr(points, to_coords=\"drive\")\n",
    "mpoints = tr(dpoints, to_coords=\"motion_space\")\n",
    "\n",
    "axs[0,0].fill(points[...,0], points[...,1])\n",
    "axs[0,1].fill(dpoints[...,0], dpoints[...,1])\n",
    "axs[0,2].fill(mpoints[...,0], mpoints[...,1])\n",
    "\n",
    "for pt, color in zip(\n",
    "    [\n",
    "        [-5, 5],\n",
    "        [-5, -5],\n",
    "        [5, -5],\n",
    "        [5, 5],\n",
    "        [0, 0]\n",
    "    ],\n",
    "    [\"red\", \"orange\", \"green\", \"purple\", \"black\"]\n",
    "):\n",
    "    dpt = tr(pt, to_coords=\"drive\")\n",
    "    mpt = tr(dpt, to_coords=\"motion_space\")\n",
    "    print(pt, dpt, mpt)\n",
    "    axs[0,0].plot(pt[0], pt[1], 'o', color=color)\n",
    "    axs[0,1].plot(dpt[..., 0], dpt[..., 1], 'o', color=color)\n",
    "    axs[0,2].plot(mpt[..., 0], mpt[..., 1], 'o', color=color)\n",
    "\n",
    "##\n",
    "\n",
    "axs[1,0].set_xlabel(\"Drive X\")\n",
    "axs[1,0].set_ylabel(\"Drive Y\")\n",
    "axs[1,1].set_xlabel(\"MSpace X\")\n",
    "axs[1,1].set_ylabel(\"MSpace Y\")\n",
    "axs[1,2].set_xlabel(\"Drive X\")\n",
    "axs[1,2].set_ylabel(\"Drive Y\")\n",
    "\n",
    "points = np.zeros((40, 2))\n",
    "points[0:10, 0] = np.linspace(-5, 5, num=10, endpoint=False)\n",
    "points[0:10, 1] = 5 * np.ones(10)\n",
    "points[10:20, 0] = 5 * np.ones(10)\n",
    "points[10:20, 1] = np.linspace(5, -5, num=10, endpoint=False)\n",
    "points[20:30, 0] = np.linspace(5, -5, num=10, endpoint=False)\n",
    "points[20:30, 1] = -5 * np.ones(10)\n",
    "points[30:40, 0] = -5 * np.ones(10)\n",
    "points[30:40, 1] = np.linspace(-5, 5, num=10, endpoint=False)\n",
    "\n",
    "mpoints = tr(points, to_coords=\"motion_space\")\n",
    "dpoints = tr(mpoints, to_coords=\"drive\")\n",
    "\n",
    "axs[1,0].fill(points[...,0], points[...,1])\n",
    "axs[1,1].fill(mpoints[...,0], mpoints[...,1])\n",
    "axs[1,2].fill(dpoints[...,0], dpoints[...,1])\n",
    "\n",
    "for pt, color in zip(\n",
    "    [\n",
    "        [-5, 5],\n",
    "        [-5, -5],\n",
    "        [5, -5],\n",
    "        [5, 5],\n",
    "        [0, 0]\n",
    "    ],\n",
    "    [\"red\", \"orange\", \"green\", \"purple\", \"black\"]\n",
    "):\n",
    "    mpt = tr(pt, to_coords=\"motion_space\")\n",
    "    dpt = tr(mpt, to_coords=\"drive\")\n",
    "    axs[1,0].plot(pt[0], pt[1], 'o', color=color)\n",
    "    axs[1,1].plot(mpt[..., 0], mpt[..., 1], 'o', color=color)\n",
    "    axs[1,2].plot(dpt[..., 0], dpt[..., 1], 'o', color=color)\n",
    "    print(f\"X = {pt[0]}  Δ = {dpt[...,0] - pt[0]} || Y = {pt[1]}  Δ = {dpt[...,1] - pt[1]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7e7e77-c0d0-4df0-bbdf-a0729792c490",
   "metadata": {},
   "source": [
    "### Test Transforming `drive -> motion space -> drive`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "094e94a9-ecbf-451a-855d-ab09e818785e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpoints = tr(points, to_coords=\"motion_space\")\n",
    "dpoints = tr(mpoints, to_coords=\"drive\")\n",
    "\n",
    "(\n",
    "    np.allclose(dpoints, points),\n",
    "    np.allclose(dpoints[...,0], points[...,0]),\n",
    "    np.allclose(dpoints[...,1], points[...,1]),\n",
    "    np.min(dpoints - points),\n",
    "    np.max(dpoints - points),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a277810b-3553-4cdf-9fc9-dc0efab86cfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "points = np.array([[5, 5], [5, 5]])\n",
    "mpoints = tr(points, to_coords=\"motion_space\")\n",
    "dpoints = tr(mpoints, to_coords=\"drive\")\n",
    "\n",
    "(\n",
    "    np.isclose(dpoints, points),\n",
    "    np.allclose(dpoints, points),\n",
    "    np.allclose(dpoints[...,0], points[...,0]),\n",
    "    np.allclose(dpoints[...,1], points[...,1]),\n",
    "    np.min(dpoints - points),\n",
    "    np.max(dpoints - points),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2cbc76-fd22-48d3-aeeb-b1650fa93f9a",
   "metadata": {},
   "source": [
    "### Test Transforming `motion space -> drive -> motion space`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8782f090-6ecb-4d25-8944-9cd1853be826",
   "metadata": {},
   "outputs": [],
   "source": [
    "dpoints = tr(points, to_coords=\"drive\")\n",
    "mpoints = tr(dpoints, to_coords=\"motion_space\")\n",
    "\n",
    "(\n",
    "    np.allclose(mpoints, points),\n",
    "    np.allclose(mpoints[...,0], points[...,0]),\n",
    "    np.allclose(mpoints[...,1], points[...,1]),\n",
    "    np.min(mpoints - points),\n",
    "    np.max(mpoints - points),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dea6120b-b00c-4828-83e8-68eff644a5e8",
   "metadata": {},
   "source": [
    "## Prototyping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0503ac38-33d2-442e-9a45-633f40bb562f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pts = [\n",
    "    [-5, 5],\n",
    "    [-5, -5],\n",
    "    [5, -5],\n",
    "    [5, 5],\n",
    "    [0, 0]\n",
    "]\n",
    "# pts = [[-5, 5]]\n",
    "\n",
    "pts = tr._condition_points(pts)\n",
    "matrix = tr.matrix(pts, to_coords=\"mspace\")\n",
    "pts = np.concatenate(\n",
    "    (pts, np.ones((pts.shape[0], 1))),\n",
    "    axis=1,\n",
    ")\n",
    "results = np.einsum(\"kmn,kn->km\", matrix, pts)[:-1,...]\n",
    "ii = 1\n",
    "# pts[ii, ...]\n",
    "(pts[ii,...], results[ii,...])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48b03845-2742-46f7-91d4-6dfa28ef3b0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix[ii, ...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a821b5-893f-4edf-93f0-2720ff5e8832",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    pts[ii, :-1],\n",
    "    tr(pts[ii, :-1], to_coords=\"mspace\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "427ec57e-aca1-4d67-82ca-c0256bdae944",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr(pts[ii, :-1], to_coords=\"mspace\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad1ea07-e6cd-4261-9baf-abdd685d35f3",
   "metadata": {},
   "source": [
    "## Testing Matrix Math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "483aa574-7a44-41b4-b2a4-3a980145a91b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pivot_to_center = 57.288\n",
    "pivot_to_drive = 134.0\n",
    "drive_polarity = np.array([1.0, 1.0])\n",
    "mspace_polarity = np.array([-1.0, 1.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287b1f92-4241-482e-97f2-f4ab1e41a9bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def matrix_to_mspace(\n",
    "    points,\n",
    "    pivot_to_center,\n",
    "    pivot_to_drive,\n",
    "    drive_polarity,\n",
    "    mspace_polarity,\n",
    "):\n",
    "    points = drive_polarity * points  # type: np.ndarray\n",
    "\n",
    "    theta = np.arctan(points[..., 1] / pivot_to_drive)\n",
    "    alpha = np.pi - theta\n",
    "\n",
    "    npoints = 1 if points.ndim == 1 else points.shape[0]\n",
    "\n",
    "    T1 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T1[..., 0, 0] = np.cos(theta)\n",
    "    T1[..., 0, 2] = -pivot_to_drive * np.cos(theta)\n",
    "    T1[..., 1, 0] = -np.sin(theta)\n",
    "    T1[..., 1, 2] = pivot_to_drive * np.sin(theta)\n",
    "    T1[..., 2, 2] = 1.0\n",
    "\n",
    "    T2 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T2[..., 0, 0] = 1.0\n",
    "    T2[..., 0, 2] = -(pivot_to_drive + pivot_to_center) * np.cos(alpha)\n",
    "    T2[..., 1, 1] = 1.0\n",
    "    T2[..., 1, 2] = -(pivot_to_drive + pivot_to_center) * np.sin(alpha)\n",
    "    T2[..., 2, 2] = 1.0\n",
    "\n",
    "    T3 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T3[..., 0, 0] = 1.0\n",
    "    T3[..., 0, 2] = -pivot_to_center\n",
    "    T3[..., 1, 1] = 1.0\n",
    "    T3[..., 2, 2] = 1.0\n",
    "    \n",
    "    # return T1, T2, T3\n",
    "    \n",
    "    T_dpolarity = np.diag(drive_polarity.tolist() + [1.0])\n",
    "    T_mpolarity = np.diag(mspace_polarity.tolist() + [1.0])\n",
    "    \n",
    "    return np.matmul(\n",
    "        T_mpolarity,\n",
    "        np.matmul(\n",
    "            T3,\n",
    "            np.matmul(\n",
    "                T2,\n",
    "                np.matmul(T1, T_dpolarity),\n",
    "            ),\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec05264-742e-40a2-8268-efebae168cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def matrix_to_drive(\n",
    "    points,\n",
    "    pivot_to_center,\n",
    "    pivot_to_drive,\n",
    "    drive_polarity,\n",
    "    mspace_polarity,\n",
    "):\n",
    "    points = mspace_polarity * points  # type: np.ndarray\n",
    "\n",
    "    # need to handle when x_L = pivot_to_center\n",
    "    # since alpha can never be 90deg we done need to worry about that case\n",
    "    alpha = np.arctan(points[..., 1] / (pivot_to_center + points[...,0]))\n",
    "\n",
    "    npoints = 1 if points.ndim == 1 else points.shape[0]\n",
    "    \n",
    "    T1 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T1[..., 0, 0] = 1.0\n",
    "    T1[..., 0, 2] = pivot_to_center\n",
    "    T1[..., 1, 1] = 1.0\n",
    "    T1[..., 2, 2] = 1.0\n",
    "\n",
    "    T2 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T2[..., 0, 0] = 1.0\n",
    "    T2[..., 0, 2] = -(pivot_to_drive + pivot_to_center) * np.cos(alpha)\n",
    "    T2[..., 1, 1] = 1.0\n",
    "    T2[..., 1, 2] = -(pivot_to_drive + pivot_to_center) * np.sin(alpha)\n",
    "    T2[..., 2, 2] = 1.0\n",
    "    \n",
    "    T3 = np.zeros((npoints, 3, 3)).squeeze()\n",
    "    T3[..., 0, 0] = 1 / np.cos(alpha)\n",
    "    T3[..., 0, 2] = pivot_to_drive\n",
    "    T3[..., 1, 2] = -pivot_to_drive * np.tan(alpha)\n",
    "    T3[..., 2, 2] = 1.0\n",
    "    \n",
    "    # return T1, T2, T3\n",
    "    \n",
    "    T_dpolarity = np.diag(drive_polarity.tolist() + [1.0])\n",
    "    T_mpolarity = np.diag(mspace_polarity.tolist() + [1.0])\n",
    "    \n",
    "    return np.matmul(\n",
    "        T_dpolarity,\n",
    "        np.matmul(\n",
    "            T3,\n",
    "            np.matmul(\n",
    "                T2,\n",
    "                np.matmul(T1, T_mpolarity),\n",
    "            ),\n",
    "        ),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49b75bb5-e2cb-4703-ab52-ada40e8ee49d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert(\n",
    "    points,\n",
    "    pivot_to_center,\n",
    "    pivot_to_drive,\n",
    "    drive_polarity,\n",
    "    mspace_polarity,\n",
    "    to_coord=\"drive\",\n",
    "):\n",
    "    if not isinstance(points, np.ndarray):\n",
    "            points = np.array(points)\n",
    "    \n",
    "    if to_coord == \"drive\":\n",
    "        matrix = matrix_to_drive(\n",
    "            points,\n",
    "            pivot_to_center=pivot_to_center,\n",
    "            pivot_to_drive=pivot_to_drive,\n",
    "            drive_polarity=drive_polarity,\n",
    "            mspace_polarity=mspace_polarity,\n",
    "        )\n",
    "    elif to_coord == \"motion_space\":\n",
    "        matrix = matrix_to_mspace(\n",
    "            points,\n",
    "            pivot_to_center=pivot_to_center,\n",
    "            pivot_to_drive=pivot_to_drive,\n",
    "            drive_polarity=drive_polarity,\n",
    "            mspace_polarity=mspace_polarity,\n",
    "        )\n",
    "    else:\n",
    "        raise ValueError\n",
    "    \n",
    "    if points.ndim == 1:\n",
    "        points = np.concatenate((points, [1]))\n",
    "        return np.matmul(matrix, points)[:2]\n",
    "\n",
    "    points = np.concatenate(\n",
    "        (points, np.ones((points.shape[0], 1))),\n",
    "        axis=1,\n",
    "    )\n",
    "    \n",
    "    return np.einsum(\"kmn,kn->km\", matrix, points)[..., :2]\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830f6bd1-cc34-44ef-8a9f-cc4b2469d097",
   "metadata": {},
   "outputs": [],
   "source": [
    "point = np.array([[0, 0], [1,2], [3,4], [-1, -1]])\n",
    "\n",
    "dpoints = convert(\n",
    "    points=point,\n",
    "    to_coord=\"drive\",\n",
    "    pivot_to_drive=pivot_to_drive,\n",
    "    pivot_to_center=pivot_to_center,\n",
    "    drive_polarity=drive_polarity,\n",
    "    mspace_polarity=mspace_polarity,\n",
    ")\n",
    "dpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1664f8b-c219-4f14-b810-9544d24af201",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpoints = convert(\n",
    "    points=dpoints,\n",
    "    to_coord=\"motion_space\",\n",
    "    pivot_to_drive=pivot_to_drive,\n",
    "    pivot_to_center=pivot_to_center,\n",
    "    drive_polarity=drive_polarity,\n",
    "    mspace_polarity=mspace_polarity,\n",
    ")\n",
    "mpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac9140f2-f9d6-4299-bc17-19433b0ddf34",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.isclose(mpoints, point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59326859-411b-41b8-9082-79d2008152a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "(mpoints - point) / point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0796106e-ecc2-495a-932f-05b59cc40257",
   "metadata": {},
   "outputs": [],
   "source": [
    "point = np.array([[0, 0], [1,2], [3,4], [-1, -1]])\n",
    "# T1, T2, T3 = matrix_to_mspace(\n",
    "#     points=point,\n",
    "#     pivot_to_center=pivot_to_center,\n",
    "#     pivot_to_drive=pivot_to_drive,\n",
    "#     drive_polarity=drive_polarity,\n",
    "#     mspace_polarity=mspace_polarity,\n",
    "# )\n",
    "T = matrix_to_mspace(\n",
    "    points=point,\n",
    "    pivot_to_center=pivot_to_center,\n",
    "    pivot_to_drive=pivot_to_drive,\n",
    "    drive_polarity=drive_polarity,\n",
    "    mspace_polarity=mspace_polarity,\n",
    ")\n",
    "TT.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a0662f-80d7-406f-9412-bbde9a8567db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (\n",
    "#     T1[1,...],\n",
    "#     T2[1,...],\n",
    "#     T3[1,...],\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7a6431-983a-432c-ba63-d13ff0894357",
   "metadata": {},
   "outputs": [],
   "source": [
    "npt = np.concatenate(\n",
    "    (\n",
    "        point,\n",
    "        np.ones((point.shape[0], 1)),\n",
    "    ),\n",
    "    axis=1,\n",
    ")\n",
    "npt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce4e703d-1357-4c9a-bdc3-c38c5ea09466",
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.matmul(TT, npt, axes=\"(k,m,n),(k,m)->(k,n)\")\n",
    "np.einsum(\"kmn,kn->km\", TT, npt)[..., :2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a67bbad7-d5c7-4ff6-836c-be94b2837187",
   "metadata": {},
   "outputs": [],
   "source": [
    "point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67082e99-d5a0-49c8-a28b-2ab0e08717fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "P = np.diag([-1, -1, 1])\n",
    "(\n",
    "    P,\n",
    "    np.linalg.inv(P),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cac88c7f-beec-4ff7-98ad-d4d2d805f952",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = np.zeros((3, 3))\n",
    "M[0,0] = 1\n",
    "M[0,2] = -50\n",
    "M[1,1] = 1\n",
    "M[2,2] = 1\n",
    "\n",
    "(\n",
    "    M,\n",
    "    np.linalg.inv(M),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377c98d5-5de0-4c83-b81a-714a7bda27b1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b8e8b58-b8e4-41ba-96a5-7f906204df89",
   "metadata": {},
   "outputs": [],
   "source": [
    "probe_axis_offset = 4.\n",
    "pivot_to_drive = 20\n",
    "pivot_to_center = 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b48262ae-75a3-499c-9d74-e2de4d0fd264",
   "metadata": {},
   "outputs": [],
   "source": [
    "points = np.array([\n",
    "    [-5, 5],\n",
    "    [-5, -5],\n",
    "    [5, -5],\n",
    "    [5, 5],\n",
    "    [0, 0],\n",
    "    [-5, 0],\n",
    "    [5, 0],\n",
    "])\n",
    "points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4493a4-4fd2-405b-b2fb-8899e6029255",
   "metadata": {},
   "outputs": [],
   "source": [
    "sine_alpha = probe_axis_offset / np.sqrt(\n",
    "    pivot_to_drive**2\n",
    "    + (-probe_axis_offset + points[..., 1])**2\n",
    ")\n",
    "alpha = np.arcsin(sine_alpha)\n",
    "np.degrees(alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b63d69-9d0b-478b-a58e-76a096de7da4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tan_beta = (-probe_axis_offset + points[..., 1]) / -pivot_to_drive\n",
    "beta = np.arctan(tan_beta)\n",
    "np.degrees(beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9eed1c-4b2a-48e8-9772-3b90920577c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = beta - alpha\n",
    "theta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a0f0cfd-af7f-4ecd-8e5b-6dbc73bdcd9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "T0 = np.zeros((points.shape[0], 3, 3)).squeeze()\n",
    "T0[..., 0, 0] = np.cos(theta)\n",
    "T0[..., 0, 2] = -pivot_to_center * (1 - np.cos(theta))\n",
    "T0[..., 1, 0] = np.sin(theta)\n",
    "T0[..., 1, 2] = pivot_to_center * np.sin(theta)\n",
    "T0[..., 2, 2] = 1.0\n",
    "T0[0,...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39cffcd6-2c05-416f-8ff1-6cce04d596a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "pts = np.concatenate(\n",
    "    (points, np.ones((points.shape[0], 1))),\n",
    "    axis=1,\n",
    ")\n",
    "mpoints = np.einsum(\"kmn,kn->km\", T0, pts)[...,:-1]\n",
    "mpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a263dd6-b393-4bbc-8359-568c8495d511",
   "metadata": {},
   "outputs": [],
   "source": [
    "tan_theta = mpoints[...,1]/(mpoints[...,0]+pivot_to_center)\n",
    "theta = -np.arctan(tan_theta)\n",
    "np.degrees(theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25bc8990-cc13-456b-a4dc-0b0bed440c2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "TI = np.zeros((points.shape[0], 3, 3)).squeeze()\n",
    "TI[..., 0, 2] = np.sqrt(mpoints[...,1]**2 +(pivot_to_center + mpoints[...,0])**2) - pivot_to_center\n",
    "TI[..., 1, 2] = pivot_to_axis * np.tan(theta) + probe_axis_offset * (1 - (1/np.cos(theta)))\n",
    "TI[..., 2, 2] = 1.0\n",
    "TI[0,...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90ac5781-d195-4019-bb2d-08dd721e7487",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpts = np.concatenate(\n",
    "    (mpoints, np.ones((points.shape[0], 1))),\n",
    "    axis=1,\n",
    ")\n",
    "pts = mpoints = np.einsum(\"kmn,kn->km\", TI, mpts)[...,:-1]\n",
    "pts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d4a1fa-cecc-4e64-97d0-b02bb8b11dd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "probe_axis_offset * (1 - (1/np.cos(theta)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6112917f-f237-46bc-a9ca-c215d2a1aef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pivot_to_axis*np.tan(theta) + probe_axis_offset * (1 - (1/np.cos(theta)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94ed0b37-ea33-4639-8f72-70008e4ac98e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pivot_to_axis*np.tan(theta) - probe_axis_offset * np.cos(theta) + probe_axis_offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82aa413b-6546-4cbc-9115-125ffcc107ee",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5853946-24e3-4286-aae6-aa48a59af280",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}