{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nearest Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Server's Data Setup\n",
    "The server owns coordinates to points of interest like restaurants and commerces. The coordinates are kept in a LookupTable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from concrete import fhe\n",
    "import numpy\n",
    "\n",
    "\n",
    "# Database of Points of Interests\n",
    "points_array = numpy.array([\n",
    "    [2, 3], [1, 5], [3, 2], [5, 2], [1, 1],\n",
    "    [9, 4], [13, 2], [14, 13], [9, 8], [8, 0],\n",
    "    [2, 10], [3, 8], [8, 12], [4, 10], [7, 7],\n",
    "])\n",
    "N_PTS = points_array.shape[0]\n",
    "points = fhe.LookupTable(points_array.flatten())\n",
    "\n",
    "\n",
    "def get_point(index):\n",
    "    return (points[2*index], points[2*index + 1])\n",
    "\n",
    "\n",
    "def all_distances(x, y):\n",
    "    xs = numpy.arange(0, 2 * N_PTS, 2)\n",
    "    ys = numpy.arange(1, 2 * N_PTS, 2)\n",
    "    a = abs(points[xs] - x)\n",
    "    b = abs(points[ys] - y)\n",
    "    return a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use swap sort to find the $K$ nearest points to a given point. However, we are interested in the indices of the elements, not just their distances. We must therefore work on tuples of index and distance, effectively implementing numpy argpartition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TLUs\n",
    "relu = fhe.univariate(lambda x: x if x > 0 else 0)\n",
    "is_positive = fhe.univariate(lambda x: 1 if x > 0 else 0)\n",
    "arg_selection = fhe.univariate(lambda x: (x-1)//2 if x % 2 else 0)  # relu packed with a flag (alternating between 0 and relu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def swap(this_idx, this_dist, that_idx, that_dist):\n",
    "    \"\"\"\n",
    "    Swaps this and that if this > that. \n",
    "    We must pass both the index and the distance for both this and that.\n",
    "\n",
    "    Returns:\n",
    "      idxmin, min, idxmax, max of this and that based on distance\n",
    "    \"\"\"\n",
    "    diff = this_dist - that_dist\n",
    "    idx = arg_selection(2 * (this_idx - that_idx) + is_positive(diff))\n",
    "    dist = relu(diff)\n",
    "\n",
    "    idx_min = this_idx - idx\n",
    "    idx_max = that_idx + idx \n",
    "    dist_min = this_dist - dist\n",
    "    dist_max = that_dist + dist\n",
    "    return fhe.array([idx_min, dist_min, idx_max, dist_max])\n",
    "\n",
    "\n",
    "@fhe.compiler({\"x\": \"encrypted\", \"y\": \"encrypted\"})\n",
    "def knn(x, y):\n",
    "    dist = all_distances(x, y)\n",
    "    idx = list(range(N_PTS))\n",
    "    for k in range(2):\n",
    "        for i in range(k+1, N_PTS):\n",
    "             idx[k], dist[k], idx[i], dist[i] = swap(idx[k], dist[k], idx[i], dist[i])\n",
    "    return fhe.array([get_point(idx[j]) for j in range(2)])\n",
    "\n",
    "\n",
    "inputset = [(4, 3), (0, 0), (15, 3), (4, 15)]\n",
    "\n",
    "circuit = knn.compile(inputset)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Client\n",
    "The client simply invokes the server's nearest neighbours circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "57.9 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 1 -n 1\n",
    "circuit.client.keys.generate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nearest(x, y):\n",
    "    ex, ey = circuit.encrypt(x, y)\n",
    "    res = circuit.run(ex, ey)   # Simulate request to the server\n",
    "    return circuit.decrypt(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21.7 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 1 -n 1\n",
    "nearest(4, 3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "zama",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}