{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import k3d\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "\n",
    "T = 1500\n",
    "N = 10\n",
    "P = 100\n",
    "mod = 50\n",
    "\n",
    "positions = [None] * N\n",
    "colors = [None] * N\n",
    "trajectories = [None] * N\n",
    "colormaps = [k3d.colormaps.matplotlib_color_maps.hot]\n",
    "\n",
    "for n in range(N):\n",
    "    positions[n] = {}\n",
    "    trajectories[n] = {}\n",
    "    bang = False\n",
    "    \n",
    "    pos = np.zeros((P, 3), dtype=np.float32)\n",
    "    velocity = np.zeros((P, 3), dtype=np.float32)\n",
    "    velocity[:, 2] = 180.0 + np.random.random(P) * 0.1    \n",
    "    \n",
    "    pos[:, 0] = np.sin(n/N * 2.0 * np.pi) * 10.0\n",
    "    pos[:, 1] = np.cos(n/N * 2.0 * np.pi) * 10.0\n",
    "    \n",
    "    cm = random.choice(colormaps)\n",
    "    cm_space = np.linspace(np.min(cm[0::4]), np.max(cm[0::4]), 2048)\n",
    "    \n",
    "    c = (np.interp(cm_space, cm[::4], cm[3::4]) * 255).astype(np.uint32) + \\\n",
    "        (np.interp(cm_space, cm[::4], cm[2::4]) * 255).astype(np.uint32) * 256 + \\\n",
    "        (np.interp(cm_space, cm[::4], cm[1::4]) * 255).astype(np.uint32) * 256 ** 2\n",
    "        \n",
    "    colors[n] = c[np.random.randint(0, c.shape[0], P)]\n",
    "    \n",
    "    for t in range(T):\n",
    "        if t > n * 30:        \n",
    "            if np.mean(velocity[:, 2] < 30.0) and not bang:\n",
    "                bang = True\n",
    "                velocity[:, 0] += np.random.normal(size=P) * 55\n",
    "                velocity[:, 1] += np.random.normal(size=P) * 55\n",
    "                velocity[:, 2] += np.random.normal(size=P) * 35\n",
    "\n",
    "            velocity[:, 2] -= 0.05\n",
    "            velocity *= 0.995                \n",
    "            pos += velocity * 0.001\n",
    "        \n",
    "        if t % mod == 0:\n",
    "            positions[n][str(t/100.0)] = pos.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in range(N):\n",
    "    db = np.array([positions[n][k] for k in positions[n].keys()])\n",
    "  \n",
    "    for t in range(T//mod):\n",
    "        anim = np.empty((P, T//mod + 1, 3), dtype=np.float32)\n",
    "        anim.fill(np.nan)\n",
    "        for p in range(P):\n",
    "            anim[p, 0:t+1, :] = db[0:t+1, p]\n",
    "            anim[p, t + 1, :] = anim[p, t, :]\n",
    "            \n",
    "        trajectories[n][str(t * mod/100.0)] = anim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "376741d20ce641d485fb90e22bd16c73",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot = k3d.plot(grid=[-100,-100,-100,100,100,100], grid_auto_fit=False, camera_auto_fit=False)\n",
    "\n",
    "for n in range(N):\n",
    "    plot += k3d.points(positions[n], colors=colors[n], shader='3d', opacity=0.85, point_size=0.5)\n",
    "    plot += k3d.line(trajectories[n], shader='simple')\n",
    "\n",
    "text_positions= {\n",
    "    '0.0': np.array([0,0,100]),\n",
    "    '4.3': np.array([0,0,100]),\n",
    "    '4.8': np.array([0,0,20])\n",
    "}\n",
    "\n",
    "text_size= {\n",
    "    '0.0': 1,\n",
    "    '4.0': 1,\n",
    "    '4.8': 80\n",
    "}\n",
    "\n",
    "plot += k3d.texture_text('Happy New Year!', text_positions, size=text_size, color=0xaaaaaa)\n",
    "plot.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.camera = [29.542883324245675,\n",
    " -25.92168776472573,\n",
    " 5.31411376357905,\n",
    " 7.077758917185657,\n",
    " -2.251048430642773,\n",
    " 15.430438268131883,\n",
    " 0.20569267725971274,\n",
    " -0.2848548297554603,\n",
    " 0.9362415545609555]\n",
    "plot.fps = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.start_auto_play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.stop_auto_play()"
   ]
  }
 ],
 "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.7.3"
  },
  "nbTranslate": {
   "displayLangs": [
    "en",
    "pl"
   ],
   "hotkey": "alt-t",
   "langInMainMenu": true,
   "sourceLang": "pl",
   "targetLang": "en",
   "useGoogleTranslate": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}