{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from e3nn import o3\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "from mpl_toolkits import mplot3d\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " X  = USV\n",
    " \n",
    " X  = USRR^{-1}V\n",
    " RV = V_tilde if R orthogonal\n",
    " \n",
    " UR^{-1}SRV if RS=SR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tetris():\n",
    "    pos = [\n",
    "        [(0, 0, 0), (0, 0, 1), (1, 0, 0), (1, 1, 0)],  # chiral_shape_1\n",
    "        [(0, 0, 0), (0, 0, 1), (1, 0, 0), (1, -1, 0)],  # chiral_shape_2\n",
    "        [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0)],  # square\n",
    "        [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3)],  # line\n",
    "        #[(0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0)],  # corner\n",
    "        [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0)],  # L\n",
    "        [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 1)],  # T\n",
    "        [(0, 0, 0), (1, 0, 0), (1, 1, 0), (2, 1, 0)],  # zigzag\n",
    "        #[(-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0)], # ?\n",
    "    ]\n",
    "    pos = torch.tensor(pos, dtype=torch.get_default_dtype())\n",
    "    # pos = pos - pos.mean(dim=1, keepdim=True)\n",
    "    print(f'This is tetris with {pos.shape[0]} blocks, described by {pos.shape[1]} {pos.shape[2]}-dim points.')\n",
    "\n",
    "    # Since chiral shapes are the mirror of one another we need an *odd* scalar to distinguish them\n",
    "    # labels = torch.tensor([\n",
    "    #     [+1, 0, 0, 0, 0, 0, 0],  # chiral_shape_1\n",
    "    #     [-1, 0, 0, 0, 0, 0, 0],  # chiral_shape_2\n",
    "    #     [0, 1, 0, 0, 0, 0, 0],  # square\n",
    "    #     [0, 0, 1, 0, 0, 0, 0],  # line\n",
    "    #     [0, 0, 0, 1, 0, 0, 0],  # corner\n",
    "    #     [0, 0, 0, 0, 1, 0, 0],  # L\n",
    "    #     [0, 0, 0, 0, 0, 1, 0],  # T\n",
    "    #     [0, 0, 0, 0, 0, 0, 1],  # zigzag\n",
    "    # ], dtype=torch.get_default_dtype())\n",
    "    # Let's do cross-entropy\n",
    "    labels = torch.arange(0, 8).to(dtype=torch.long)\n",
    "\n",
    "    # apply random rotation\n",
    "    rotation = o3.rand_matrix(len(pos))\n",
    "    rotated_positions = torch.einsum('zij,zaj->zai', o3.rand_matrix(len(pos)), pos)\n",
    "\n",
    "    # put in torch_geometric format\n",
    "    # dataset = [Data(pos=pos) for pos in pos]\n",
    "    # data = next(iter(DataLoader(dataset, batch_size=len(dataset))))\n",
    "\n",
    "    return rotated_positions, labels, rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos, labels, rotation = tetris()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Align Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pos1, labels, rotation = tetris()\n",
    "pos2, labels, rotation = tetris()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pos1 = pos1.permute(0, 2, 1)\n",
    "# pos2 = pos2.permute(0, 2, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def align(positions):\n",
    "    X = positions - positions.mean(dim=1, keepdim=True)\n",
    "    #C = X.transpose(1, 2).matmul(X)\n",
    "    #ew, ev = torch.symeig(C, eigenvectors=True)\n",
    "    #print(torch.round(100 * ew) / 100)\n",
    "    #    return X.matmul(ev.transpose(1, 2))\n",
    "    \n",
    "    U, S, V = torch.linalg.svd(X, full_matrices=False)\n",
    "    #print(torch.round(S * 1000) / 1000)\n",
    "    repeats = ((torch.isclose(S[:,0], S[:, 1]) | torch.isclose(S[:,1], S[:,0])).nonzero())\n",
    "    if len(repeats) > 0:\n",
    "        print(repeats)\n",
    "    print(S)\n",
    "    d = torch.sign(U[:, 0, :]) # get the first significant sign of u # S is ordered\n",
    "    # d = torch.nn.functional.hardshrink(U[:, :,:], 0).max.sign() # get the first significant sign of u\n",
    "    # most_significant_digit = U.abs().argmax(dim=1, keepdim=True) # get the most significant sign of u\n",
    "    # d = torch.sign(U.gather(dim=1, index=most_significant_digit).squeeze())\n",
    "    # print(U[:, 0, :])\n",
    "    # print(torch.nn.functional.hardshrink(U[:, :,:], 10).max(dim=1)[0])\n",
    "    #print(d)\n",
    "    D = torch.diag_embed(d.float())\n",
    "    # print(D, D.shape)\n",
    "\n",
    "    #return U.matmul(torch.diag_embed(S)).matmul(D)\n",
    "    return X.matmul(V.transpose(1, 2).matmul(D))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "align1 = align(pos1)\n",
    "align2 = align(pos2)\n",
    "\n",
    "#(align1 - align2).abs() \n",
    "(align1 - align2).abs() * ((align1 - align2).abs()>1e-6)\n",
    "#(align1.abs() - align2.abs()).abs() * ((align1.abs() - align2.abs()).abs()>1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "align1[4], align2[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for thing in range(4):\n",
    "    fig = plt.figure(figsize=(2,2))\n",
    "    ax = plt.axes(projection='3d')\n",
    "\n",
    "    # Data for three-dimensional scattered points\n",
    "    shape = 0\n",
    "\n",
    "    zdata = X_out[thing][shape, :, 0].numpy()\n",
    "    xdata = X_out[thing][shape, :, 1].numpy()\n",
    "    ydata = X_out[thing][shape, :, 2].numpy()\n",
    "    ax.plot3D(xdata, ydata, zdata, color='b');\n",
    "    ax.set_xlim3d(-2, 2)\n",
    "    ax.set_ylim3d(-2, 2)\n",
    "    ax.set_zlim3d(-2, 2)\n",
    "    \n",
    "\n",
    "for thing in range(4):\n",
    "    fig = plt.figure(figsize=(2,2))\n",
    "    ax = plt.axes(projection='3d')\n",
    "\n",
    "    # Data for three-dimensional scattered points\n",
    "    shape = 0\n",
    "\n",
    "    zdata = X_out[thing][shape, :, 0].numpy()\n",
    "    xdata = X_out[thing][shape, :, 1].numpy()\n",
    "    ydata = X_out[thing][shape, :, 2].numpy()\n",
    "    ax.plot3D(xdata, ydata, zdata, color='b');\n",
    "    ax.set_xlim3d(-2, 2)\n",
    "    ax.set_ylim3d(-2, 2)\n",
    "    ax.set_zlim3d(-2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.dist(align1, align2)"
   ]
  }
 ],
 "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
