{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "4e107ea1-d6d3-4bc2-b2fc-32d6f283447e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import trimesh\n",
    "import pysdf\n",
    "import k3d\n",
    "import numpy as np\n",
    "import scipy\n",
    "from ortools.init.python import init\n",
    "from ortools.linear_solver import pywraplp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "1f25c5c1-97e5-4d27-8a5c-5b50d2c931bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: pysdf chokes on FreeCad .obj files imported through trimesh.\n",
    "\n",
    "mesh = trimesh.load('/tmp/feedthrough1_cut2-PrimaryMeshbracket.stl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "id": "efed4eb5-0779-4be1-9849-9816dfd81c5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jaseg/proj/ihsm-secondary-mesh/venv/lib/python3.12/site-packages/traittypes/traittypes.py:97: UserWarning: Given trait value dtype \"int64\" does not match required type \"float32\". A coerced copy has been created.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bda3e2161d89497ca8ee788334884892",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "line = np.array([\n",
    "    (3, 0, -13),\n",
    "    (23, 0, -13),\n",
    "    (23, 0, -2),\n",
    "    (40, 0, -2),\n",
    "    (40, 0, -13),\n",
    "    (55, 0, -13),\n",
    "    (55, 0, -2),\n",
    "    (75, 0, -2),\n",
    "    (85, 0, -8)])\n",
    "\n",
    "plot = k3d.plot()\n",
    "kmesh = k3d.mesh(mesh.vertices, mesh.faces, color=0xdc8add, opacity=0.8)\n",
    "plot += kmesh\n",
    "plot += k3d.label('Start', position=line[0])\n",
    "plot += k3d.label('End', position=line[-1])\n",
    "plot += k3d.line(line)\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "4215b9e5-4533-451a-94be-0fc62ade9b34",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = pysdf.SDF(mesh.vertices, mesh.faces)\n",
    "\n",
    "xs = np.linspace(-10, 90, 2000, dtype=np.float32)\n",
    "ys = np.linspace(-10, 10, 40, dtype=np.float32)\n",
    "zs = np.linspace(-14, 0, 32, dtype=np.float32)\n",
    "coords = np.array(np.meshgrid(xs, ys, zs, indexing='ij')).T.reshape(-1, 3)\n",
    "sdf_vol = f(coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "6b955c31-639d-49a8-983d-744bf9648a38",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3fb11fc88976448990c5ee91669cfa77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot = k3d.plot()\n",
    "plot += k3d.points(positions=coords,\n",
    "                   attribute=sdf_vol,\n",
    "                  point_size=0.1,\n",
    "                  opacity=0.2,\n",
    "                  shader='flat')\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "id": "d3a63bbe-a827-4fe9-8806-50081d844702",
   "metadata": {},
   "outputs": [],
   "source": [
    "def space_to_cloud(x, y, z):\n",
    "    return min(len(xs), max(0, round((x+10)/100 * 2000))),\\\n",
    "           min(len(ys), max(0, round((y+10)/20 * 40))),\\\n",
    "           min(len(zs), max(0, round((z+14)/14 * 32)))\n",
    "    \n",
    "def cloud_to_space(x, y, z):\n",
    "    return np.vstack([x/2000 * 100 - 10,\n",
    "           y/40 * 20 - 10,\n",
    "           z/32 * 14 - 14]).T\n",
    "    \n",
    "line_interp = np.vstack([np.interp(np.linspace(0, 1, 100), np.linspace(0, 1, len(line)), np.array(line)[:,i]) for i in range(3)]).T\n",
    "line_cloud = np.array([space_to_cloud(x, y, z) for x, y, z in line_interp], dtype=int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "id": "ef4b7b64-70ba-4e26-90c1-240be0bb48ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Objective: 34.75516\n"
     ]
    }
   ],
   "source": [
    "def objective(line_cloud):\n",
    "    return np.sqrt((sdf_vol[line_cloud]**2).sum())\n",
    "    \n",
    "print('Objective:', objective(line_cloud))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "id": "a497cb6b-32dc-4b67-8a84-6a2a99e68fdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e55c322c205d42b2ae2cc0d4354964ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sdf_interp = scipy.interpolate.RegularGridInterpolator((xs, ys, zs), sdf_vol.reshape(len(zs), len(ys), len(xs)).swapaxes(0,2))\n",
    "\n",
    "plot = k3d.plot()\n",
    "sel = (coords[:,1]%-8 > -0.2)\n",
    "plot += k3d.points(positions=coords[sel],\n",
    "                   attribute=sdf_interp(coords[sel]),\n",
    "                  point_size=0.6,\n",
    "                  opacity=0.2,\n",
    "                  shader='flat')\n",
    "plot += k3d.line([cloud_to_space(x, y, z) for x, y, z in line_cloud])\n",
    "\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "b523f8fd-9119-4c9a-98bb-20f583182509",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.002690871616336777 0.009143656777662025\n",
      "0.011834528393998803\n"
     ]
    }
   ],
   "source": [
    "def objective_space(idx, debug=False):\n",
    "    idx = idx.reshape(-1, 3)\n",
    "    idx = np.vstack([line[:1], idx, line[-1:]])\n",
    "    deltas = idx[:-1] - idx[1:]\n",
    "    dists = np.sqrt(deltas[:,0]**2 + deltas[:,1]**2 + deltas[:,2]**2)\n",
    "    \n",
    "    sdf_vals = sdf_interp(idx)\n",
    "    sdf_comp = 1/(0.00001 + np.abs(sdf_vals + .1)).sum()\n",
    "    dist_comp = 0.1* ((dists - (dists.sum()/len(dists)))**2).mean()\n",
    "\n",
    "    if debug:\n",
    "        print(sdf_comp, dist_comp)\n",
    "\n",
    "    return sdf_comp + dist_comp\n",
    "\n",
    "print(objective_space(line_interp[1:-1], debug=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "id": "a31cd866-093b-46f4-9c50-c83e0b84d3d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  message: CONVERGENCE: REL_REDUCTION_OF_F_<=_FACTR*EPSMCH\n",
      "  success: True\n",
      "   status: 0\n",
      "      fun: 0.0020290044959345027\n",
      "        x: [ 4.058e+00  3.529e-02 ...  1.216e-04 -7.654e+00]\n",
      "      nit: 233\n",
      "      jac: [-8.357e-06 -2.846e-06 ...  4.337e-11  1.758e-07]\n",
      "     nfev: 86730\n",
      "     njev: 294\n",
      " hess_inv: <294x294 LbfgsInvHessProduct with dtype=float64>\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4aec3152a484d59b060731442352a7b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bounds = np.tile([[-10, 90], [-10, 10], [-14, 0]], (len(line_interp)-2, 1))\n",
    "res = scipy.optimize.minimize(objective_space, x0=line_interp[1:-1].flatten(), tol=1e-8,\n",
    "                              bounds=bounds, method='L-BFGS-B', options=dict(maxiter=100000, maxfun=100000, maxls=40))\n",
    "print(repr(res))\n",
    "assert res.success\n",
    "\n",
    "plot = k3d.plot()\n",
    "plot += k3d.mesh(mesh.vertices, mesh.faces, color=0xdc8add, opacity=0.8)\n",
    "plot += k3d.line(res.x.reshape(-1, 3))\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ada64018-2fbf-44ce-8880-3680b0195894",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ihsm-secmesh",
   "language": "python",
   "name": "ihsm-secmesh"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}