{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eaed683c-c6f1-45e4-aaee-ae4e57209f5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy as scp\n",
    "from scipy.spatial import Voronoi, voronoi_plot_2d\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "efd84ecd-9fb3-4d2f-9f7a-47cd1ff09eea",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Lattice:\n",
    "    def __init__(self, *vecs):\n",
    "        # if the vecs were put in an iterable\n",
    "        if len(vecs) == 1:\n",
    "            vecs = vecs[0]\n",
    "        if len(vecs)  == 3:\n",
    "            pass\n",
    "        elif len(vecs) == 2:\n",
    "            pass\n",
    "        else: raise ValueError(\"Vecs must contain either 2 or 3 vectors\")\n",
    "        self.dim = len(vecs)\n",
    "        self.vecs = list(vecs)\n",
    "        for i, v in enumerate(self.vecs):\n",
    "            if type(v) != np.ndarray:\n",
    "                self.vecs[i] = np.array(v)\n",
    "            if self.vecs[i].shape != (self.dim,):\n",
    "                raise ValueError(f\"Got {self.dim} vectors, therefore all vectors must be {self.dim} dimensional but vector {i+1} has shape {self.vecs[i].shape}\")\n",
    "        self.vecs = np.array(self.vecs)\n",
    "        self.vec_lengths = np.array([np.linalg.norm(v) for v in self.vecs])\n",
    "        self.center = np.zeros(self.dim)\n",
    "\n",
    "    def get_point(self, *ns):\n",
    "        if len(ns) != len(self.vecs): raise ValueError(f\"Requires one index for each lattice vector {len(self.vecs)}, but got only {ns}\")\n",
    "        point = self.center.copy()\n",
    "        for i, n in enumerate(ns):\n",
    "            point += n * self.vecs[i]\n",
    "        return point\n",
    "\n",
    "    \n",
    "    def get_points_around_center(self, n):\n",
    "        import itertools\n",
    "        ns = [i for i in range(-i, i+1)]\n",
    "        for n in itertools.product(*[])\n",
    "        for i in range(self.dim):\n",
    "            \n",
    "            \n",
    "        for i in range(-n, n+1):\n",
    "            for j "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2855d08e-70d8-4ef7-ba19-2c06316caf15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<scipy.spatial._qhull.Voronoi object at 0x7e3efffb4a70>\n",
      "1.7320508075688774 [0.70710678 1.58113883 0.70710678 0.70710678 1.58113883 0.70710678]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24624 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21868 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sphere_point = lambda rad: np.array([np.cos(rad), np.sin(rad)])\n",
    "\n",
    "square_lattice = Lattice([1, 0], [0, 1])\n",
    "tilted_lattice = Lattice([1, 1], [0, 1])\n",
    "honeycomb_lattice = Lattice(sphere_point(0), sphere_point(np.pi * 2/3))\n",
    "\n",
    "def get_elementary_cell(lattice: Lattice):\n",
    "    points = []\n",
    "    for n1 in range(-1, 2):\n",
    "        for n2 in range(-1, 2):\n",
    "            p = lattice.get_point(n1, n2)\n",
    "            p_len = np.linalg.norm(p)\n",
    "            # if (p_len <= lattice.vec_lengths).any():\n",
    "            points.append(p)\n",
    "    return points\n",
    "\n",
    "def get_n_points_around_center(n)\n",
    "\n",
    "points = get_elementary_cell(tilted_lattice)\n",
    "vor = Voronoi(points)\n",
    "\n",
    "print(vor)\n",
    "fig = voronoi_plot_2d(vor)\n",
    "\n",
    "def get_orthogonal_2D(vec):\n",
    "    return np.array((vec[1], -vec[0]))\n",
    "def plot_voronoi(lattice, v: Voronoi, subplot_kw={}):\n",
    "    fig, ax = plt.subplots(**subplot_kw) \n",
    "    lattice_points = v.points\n",
    "    lattice_vec_norm = np.sqrt(np.sum(lattice.vec_lengths**2))\n",
    "    print(lattice_vec_norm, np.linalg.norm(v.vertices, axis=1))\n",
    "    cell_vertices = v.vertices[np.linalg.norm(v.vertices, axis=1) <= 0.5 * lattice_vec_norm]\n",
    "    \n",
    "    x, y = zip(*lattice_points)\n",
    "    ax.scatter(x, y, color=\"blue\")\n",
    "    x, y = zip(*cell_vertices)\n",
    "    arrowprops = dict(arrowstyle=\"-|>,head_width=0.4,head_length=0.8\", color=\"black\", shrinkA=0,shrinkB=0)\n",
    "    for i, vec in enumerate(lattice.vecs):\n",
    "        ax.annotate(f\"\", xy=lattice.vecs[i], xytext=lattice.center, arrowprops=arrowprops)\n",
    "        # add name of vector at a perpendicular offset starting at half length\n",
    "        ax.annotate(r\"$\\vec{a}_\"+f\"{i+1}$\", xy=0.7*lattice.vecs[i], xytext=0.7*lattice.vecs[i] + 0.06*get_orthogonal_2D(lattice.vecs[i]))\n",
    "        \n",
    "    ax.scatter(x, y, color=\"orange\")\n",
    "    ax.fill(x, y, color=\"#4444\")\n",
    "    return fig\n",
    "    \n",
    "\n",
    "plot_voronoi(tilted_lattice, vor)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2faef81-5f2a-4950-b986-e891ddfa7da8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6b6ccf4-34c7-419d-8bcc-f3ce2870ba25",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda",
   "language": "python",
   "name": "conda"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
