{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>\n",
       ".cell-output-ipywidget-background {\n",
       "    background-color: transparent !important;\n",
       "}\n",
       ":root {\n",
       "    --jp-widgets-color: var(--vscode-editor-foreground);\n",
       "    --jp-widgets-font-size: var(--vscode-editor-font-size);\n",
       "}  \n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "<style>\n",
    ".cell-output-ipywidget-background {\n",
    "    background-color: transparent !important;\n",
    "}\n",
    ":root {\n",
    "    --jp-widgets-color: var(--vscode-editor-foreground);\n",
    "    --jp-widgets-font-size: var(--vscode-editor-font-size);\n",
    "}  \n",
    "</style>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import pyaurorax\n",
    "\n",
    "aurorax = pyaurorax.PyAuroraX()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3eb5632854fb4f5488fe9f90e8bcacd8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading THEMIS_ASI_RAW files:   0%|          | 0.00/128M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# download an hour of THEMIS ASI data\n",
    "dataset_name = \"THEMIS_ASI_RAW\"\n",
    "start_dt = datetime.datetime(2021, 11, 4, 9, 0)\n",
    "end_dt = datetime.datetime(2021, 11, 4, 9, 59)\n",
    "site_uid = \"atha\"\n",
    "r = aurorax.data.ucalgary.download(dataset_name, start_dt, end_dt, site_uid=site_uid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in the hour of data\n",
    "data = aurorax.data.ucalgary.read(r.dataset, r.filenames, n_parallel=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up working with the tools by just making a shorter name for our future calls\n",
    "at = aurorax.tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# scale all the images\n",
    "#\n",
    "# NOTE: you can scale all images or just one image\n",
    "images_scaled = at.scale_intensity(data.data, min=1000, max=10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "330cd24cba294a76a53eb26c5102206d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating frame files:   0%|          | 0/1200 [00:00<?, ?frames/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# we're going to process all frames from the hour. We'll use the display() function\n",
    "# and have it return the figure, then we'll add some timestamp and other text\n",
    "# to our liking.\n",
    "#\n",
    "# we'll utilize multiprocessing and tqdm's progress bar functionality to do this\n",
    "# using several worker processes\n",
    "\n",
    "import os\n",
    "import platform\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.contrib.concurrent import process_map as tqdm_process_map\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "\n",
    "def process_frame(i):\n",
    "    fig, ax = at.display(images_scaled[:, :, i], cmap=\"gray\", returnfig=True)\n",
    "    ax.text(5, 240, \"THEMIS ASI\", color=\"white\", size=14)\n",
    "    ax.text(5, 225, \"ATHA\", color=\"white\", size=14)\n",
    "    ax.text(145, 8, data.timestamp[i].strftime(\"%Y-%m-%d %H:%M:%S UTC\"), color=\"white\", size=11)\n",
    "    filename = \"movie_frames/%s_atha_themis.png\" % (data.timestamp[i].strftime(\"%Y%m%d_%H%M%S\"))\n",
    "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
    "    plt.savefig(filename, bbox_inches=\"tight\")\n",
    "    plt.close()\n",
    "    return filename\n",
    "\n",
    "\n",
    "if (platform.system() == \"Windows\"):\n",
    "    # pre-process frames serially\n",
    "    #\n",
    "    # NOTE: multiprocessing on Windows from within a notebook is not too easy, so we'll\n",
    "    # just do this serially.\n",
    "    frame_filename_list = []\n",
    "    for i in tqdm(range(0, images_scaled.shape[-1]), total=images_scaled.shape[-1], desc=\"Generating frame files: \", unit=\"frames\"):\n",
    "        frame_filename_list.append(process_frame(i))\n",
    "else:\n",
    "    frame_filename_list = tqdm_process_map(\n",
    "        process_frame,\n",
    "        range(0, images_scaled.shape[-1]),\n",
    "        max_workers=5,\n",
    "        chunksize=1,\n",
    "        desc=\"Generating frame files: \",\n",
    "        unit=\"frames\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "676c8ca40cb241b38999200ff500662e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Reading files:   0%|          | 0/1200 [00:00<?, ?files/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2cdd0d00797b4dfe9323fb215f3df5f3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Encoding frames:   0%|          | 0/1200 [00:00<?, ?frames/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# now that we have our frames, we'll generate a movie using them\n",
    "at.movie(frame_filename_list, \"test.mp4\", n_parallel=5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyaurorax39",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}