{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modified from: https://twitter.com/balemkevin/status/1271084879518859267\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xarray\n",
    "from argopy import DataFetcher as ArgoDataFetcher\n",
    "\n",
    "lon_min, lon_max = -136, -120\n",
    "lat_min, lat_max = 41.78, 52.24\n",
    "\n",
    "\n",
    "argo_loader = ArgoDataFetcher(cache=False)\n",
    "\n",
    "surface = 0, 10\n",
    "deep = 975, 1025\n",
    "time_span = \"2009-12\", \"2020-01\"\n",
    "\n",
    "df1 = argo_loader.region(\n",
    "    [lon_min, lon_max, lat_min, lat_max, *surface, *time_span]\n",
    ").to_xarray()\n",
    "df2 = argo_loader.region(\n",
    "    [lon_min, lon_max, lat_min, lat_max, *deep, *time_span]\n",
    ").to_xarray()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# weeks\n",
    "daterange = np.arange(\"2010-01-01\", \"2020-01-03\", dtype=\"datetime64[7D]\")\n",
    "\n",
    "# middle of the week\n",
    "dayoftheyear = pd.DatetimeIndex(\n",
    "    np.array(daterange, dtype=\"datetime64[D]\") + 3\n",
    ").dayofyear\n",
    "activeyear = pd.DatetimeIndex(np.array(daterange, dtype=\"datetime64[D]\") + 3).year\n",
    "\n",
    "tsurf = np.zeros(len(daterange))\n",
    "t1000 = np.zeros(len(daterange))\n",
    "\n",
    "for k in range(len(daterange)):\n",
    "    i1 = (df1[\"TIME\"] >= daterange[k]) & (df1[\"TIME\"] < daterange[k] + 7)\n",
    "    i2 = (df2[\"TIME\"] >= daterange[k]) & (df2[\"TIME\"] < daterange[k] + 7)\n",
    "    tsurf[k] = df1.where(i1, drop=True)[\"TEMP\"].mean().values\n",
    "    t1000[k] = df2.where(i2, drop=True)[\"TEMP\"].mean().values\n",
    "\n",
    "d = {\n",
    "    \"date\": np.array(daterange, dtype=\"datetime64[D]\"),\n",
    "    \"tsurf\": tsurf,\n",
    "    \"t1000\": t1000,\n",
    "}\n",
    "ndf = pd.DataFrame(data=d)\n",
    "ndf.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.rcParams[\"xtick.major.pad\"] = \"17\"\n",
    "plt.rcParams[\"axes.axisbelow\"] = False\n",
    "matplotlib.rc(\"axes\", edgecolor=\"w\")\n",
    "from IPython.display import HTML\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "big_angle = 360 / 12\n",
    "date_angle = ((360 / 365) * dayoftheyear) * np.pi / 180\n",
    "inner = 2\n",
    "outer = 19\n",
    "ocean_color = [\"#ff7f50\", \"#004752\"]\n",
    "\n",
    "\n",
    "def dress_axes(ax):\n",
    "    ax.set_facecolor(\"w\")\n",
    "    ax.set_theta_zero_location(\"N\")\n",
    "    ax.set_theta_direction(-1)\n",
    "    middles = np.arange(big_angle / 2, 360, big_angle) * np.pi / 180\n",
    "    ax.set_xticks(middles)\n",
    "    ax.set_xticklabels(\n",
    "        [\n",
    "            \"January\",\n",
    "            \"February\",\n",
    "            \"March\",\n",
    "            \"April\",\n",
    "            \"May\",\n",
    "            \"June\",\n",
    "            \"July\",\n",
    "            \"August\",\n",
    "            \"September\",\n",
    "            \"October\",\n",
    "            \"November\",\n",
    "            \"December\",\n",
    "        ]\n",
    "    )\n",
    "    ax.set_yticks([2, 8, 19])\n",
    "    ax.set_yticklabels([\"2°C\", \"8°C\", \"19°C\"])\n",
    "    ax.set_rlabel_position(359)\n",
    "    ax.tick_params(axis=\"both\", color=\"w\")\n",
    "    plt.grid(None, axis=\"x\")\n",
    "    plt.grid(axis=\"y\", color=\"w\", linestyle=\":\", linewidth=1)\n",
    "    bars = ax.bar(\n",
    "        middles,\n",
    "        outer,\n",
    "        width=big_angle * np.pi / 180,\n",
    "        bottom=inner,\n",
    "        color=\"lightgray\",\n",
    "        edgecolor=\"w\",\n",
    "        zorder=0,\n",
    "    )\n",
    "    plt.ylim([2, outer])\n",
    "    legend_elements = [\n",
    "        Line2D(\n",
    "            [0],\n",
    "            [0],\n",
    "            marker=\"o\",\n",
    "            color=\"w\",\n",
    "            label=\"Surface\",\n",
    "            markerfacecolor=ocean_color[0],\n",
    "            markersize=15,\n",
    "        ),\n",
    "        Line2D(\n",
    "            [0],\n",
    "            [0],\n",
    "            marker=\"o\",\n",
    "            color=\"w\",\n",
    "            label=\"1000m\",\n",
    "            markerfacecolor=ocean_color[1],\n",
    "            markersize=15,\n",
    "        ),\n",
    "    ]\n",
    "    ax.legend(handles=legend_elements, loc=\"center\", fontsize=13, frameon=False)\n",
    "    plt.suptitle(\n",
    "        \"Temperature from Argo profiles\",\n",
    "        fontsize=16,\n",
    "        horizontalalignment=\"center\",\n",
    "    )\n",
    "\n",
    "\n",
    "def update(i):\n",
    "    ax.cla()\n",
    "    dress_axes(ax)\n",
    "    i0 = np.max([i - 51, 0])\n",
    "\n",
    "    ax.plot(\n",
    "        date_angle[i0 : i + 1],\n",
    "        ndf[\"tsurf\"][i0 : i + 1],\n",
    "        \"-\",\n",
    "        color=ocean_color[0],\n",
    "        alpha=1.0,\n",
    "        linewidth=5,\n",
    "    )\n",
    "    ax.plot(\n",
    "        date_angle[0 : i + 1],\n",
    "        ndf[\"tsurf\"][0 : i + 1],\n",
    "        \"-\",\n",
    "        color=ocean_color[0],\n",
    "        linewidth=0.7,\n",
    "    )\n",
    "\n",
    "    ax.plot(\n",
    "        date_angle[i0 : i + 1],\n",
    "        ndf[\"t1000\"][i0 : i + 1],\n",
    "        \"-\",\n",
    "        color=ocean_color[1],\n",
    "        alpha=1.0,\n",
    "        linewidth=5,\n",
    "    )\n",
    "    ax.plot(\n",
    "        date_angle[0 : i + 1],\n",
    "        ndf[\"t1000\"][0 : i + 1],\n",
    "        \"-\",\n",
    "        color=ocean_color[1],\n",
    "        linewidth=0.7,\n",
    "    )\n",
    "\n",
    "    ax.plot([date_angle[i], date_angle[i]], [inner, outer], \"k-\", linewidth=0.5)\n",
    "    plt.title(str(activeyear[i]), fontsize=16, horizontalalignment=\"center\")\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, polar=True)\n",
    "\n",
    "dress_axes(ax)\n",
    "update(322)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "anim = FuncAnimation(fig, update, interval=40, frames=len(daterange) - 1, repeat=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML(anim.to_html5_video())"
   ]
  }
 ],
 "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}