{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7772f25c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyvista as pv\n",
    "import pvxarray\n",
    "import xarray as xr\n",
    "\n",
    "pv.set_plot_theme(\"document\")\n",
    "pv.set_jupyter_backend(\"server\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21ff7e10-b4b6-4238-8958-e43e0d492ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Temporary manual mode override; waiting for pyvista release including changes from https://github.com/pyvista/pyvista/pull/5798\n",
    "\n",
    "pv.global_theme.trame.server_proxy_enabled = False\n",
    "pv.global_theme.trame.jupyter_extension_enabled = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ea7e2db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(\"data/all_flashes.nc\")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96db61f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# event_parent_flash_id , flash_event_count (min 10)\n",
    "ds.flash_id[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8dc6b0f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = plt.hist(ds.event_z, bins=100, range=(400, 15000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b2dbda",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = plt.hist(ds.event_parent_flash_id, bins=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95500c41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds.event_parent_flash_id[ds.event_parent_flash_id == 2000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf5373d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "points = np.c_[ds.event_x, ds.event_y, ds.event_z]\n",
    "\n",
    "t = pd.to_datetime(ds.event_time).astype(np.int64)\n",
    "\n",
    "\n",
    "def get_flash(flash_index):\n",
    "    flash_id = ds.event_parent_flash_id[flash_index]\n",
    "    event_mask = ds.event_parent_flash_id == flash_id\n",
    "\n",
    "    pc = pv.PolyData(points[event_mask])\n",
    "    pc[\"event_time\"] = t[event_mask]  # da.data\n",
    "    pc[\"event_parent_flash_id\"] = ds.event_parent_flash_id[event_mask]\n",
    "\n",
    "    pc = pc.elevation().threshold((400, 15000), scalars=\"Elevation\").extract_geometry()\n",
    "    return pc\n",
    "\n",
    "\n",
    "pc = pv.PolyData(points)\n",
    "pc[\"event_time\"] = t.values\n",
    "pc[\"event_parent_flash_id\"] = ds.event_parent_flash_id.values\n",
    "\n",
    "full_pc = pc.elevation().threshold((400, 15000), scalars=\"Elevation\")\n",
    "box = full_pc.outline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad02db92-229c-41e9-aeae-2f9754d2884f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# full_pc.extract_surface().save('lightning.vtp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36564f65",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pc = get_flash(2000)\n",
    "\n",
    "pl = pv.Plotter()\n",
    "actor = pl.add_mesh(\n",
    "    pc,\n",
    "    scalars=\"event_time\",\n",
    "    cmap=\"plasma\",\n",
    "    point_size=10,\n",
    "    ambient=0.5,\n",
    ")\n",
    "pl.set_background(\"grey\")\n",
    "pl.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7cf7c0f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "import pyvista as pv\n",
    "from tqdm import tqdm\n",
    "\n",
    "from pvxarray.vtk_source import PyVistaXarraySource\n",
    "\n",
    "\n",
    "def time_controls(plotter: pv.BasePlotter, continuous_update=False, step=1):\n",
    "    def update_time_index(time_index):\n",
    "        ...\n",
    "        plotter.add_mesh(\n",
    "            get_flash(time_index),\n",
    "            scalars=\"event_time\",\n",
    "            cmap=\"plasma\",\n",
    "            point_size=10,\n",
    "            ambient=0.5,\n",
    "            render_points_as_spheres=True,\n",
    "            name=\"flash\",\n",
    "        )\n",
    "\n",
    "        plotter.render()\n",
    "\n",
    "    tmax = len(ds.event_time)\n",
    "\n",
    "    def set_time(change):\n",
    "        value = change[\"new\"]\n",
    "        if value < 0:\n",
    "            value = 0\n",
    "        if value >= tmax:\n",
    "            value = tmax - 1\n",
    "        update_time_index(value)\n",
    "\n",
    "    play = widgets.Play(\n",
    "        value=0,\n",
    "        min=0,\n",
    "        max=tmax,\n",
    "        step=step,\n",
    "        description=\"Time Index\",\n",
    "    )\n",
    "    play.observe(set_time, \"value\")\n",
    "\n",
    "    slider = widgets.IntSlider(min=0, max=tmax, continuous_update=continuous_update)\n",
    "    widgets.jslink((play, \"value\"), (slider, \"value\"))\n",
    "    return widgets.HBox([play, slider])\n",
    "\n",
    "\n",
    "pl = pv.Plotter()\n",
    "actor = pl.add_mesh(\n",
    "    get_flash(0),\n",
    "    scalars=\"event_time\",\n",
    "    cmap=\"plasma\",\n",
    "    point_size=10,\n",
    "    ambient=0.5,\n",
    "    name=\"flash\",\n",
    "    render_points_as_spheres=True,\n",
    ")\n",
    "pl.add_mesh(box, color=\"k\")\n",
    "pl.set_background(\"grey\")\n",
    "pl.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "355a4a40",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "time_controls(pl, continuous_update=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}