{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "matplotlib.use('nbagg')\n",
    "\n",
    "import pandas as pd\n",
    "import sys\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "import zarr\n",
    "import warnings\n",
    "from textwrap import wrap\n",
    "from mpl_toolkits import mplot3d\n",
    "import random\n",
    "from scipy.ndimage.filters import uniform_filter\n",
    "from scipy.ndimage.measurements import variance\n",
    "from scipy.ndimage import label\n",
    "from scipy.ndimage.morphology import binary_closing\n",
    "from skimage.filters import gaussian, threshold_otsu\n",
    "from skimage import measure\n",
    "\n",
    "from re import split\n",
    "\n",
    "import dask\n",
    "from ipywidgets import interact\n",
    "from dask_jobqueue import SLURMCluster\n",
    "from dask.distributed import Client, progress, wait, fire_and_forget\n",
    "from dask import delayed, compute, visualize\n",
    "from dask_image.ndfilters import uniform_filter as uf\n",
    "from dask_image.ndfilters import gaussian_filter\n",
    "from dask_image.ndmeasure import variance as varian\n",
    "import dask_image.ndmeasure as da_measure\n",
    "import dask.array as da\n",
    "\n",
    "import time\n",
    "warnings.filterwarnings('ignore')\n",
    "sys.path.insert(1, f\"{os.path.abspath(os.path.join(os.path.abspath(''), '../'))}\")\n",
    "from src.utils import time_3d, get_pars_from_ini\n",
    "location = split(', |_|-|!', os.popen('hostname').read())[0].replace(\"\\n\", \"\")\n",
    "path_data = get_pars_from_ini(campaign='loc')[location]['path_data']\n",
    "path_proj = get_pars_from_ini(campaign='loc')[location]['path_proj']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @dask.delayed\n",
    "def multiple_plot(ds):\n",
    "    if sum(ds.zhh14.shape) > 0:\n",
    "        plt.close('all')\n",
    "        fig, axs = plt.subplots(3,2,  figsize=(11,20), sharey=True)\n",
    "        var = {'zhh14': {'cmp': 'jet', 'vim':0, 'vmax':40, 'name': 'Reflectivity'}, \n",
    "               'zhh14SP': {'cmp': 'jet', 'vim':0, 'vmax':40, 'name': 'Reflectivity'}, \n",
    "               'vel14': {'cmp': 'hsv', 'vim':-15, 'vmax':15, 'name': 'Velocity'}, \n",
    "               'vel14SP': {'cmp': 'hsv', 'vim':-15, 'vmax':15, 'name': 'Velocity'}, \n",
    "               'ldrhh14': {'cmp': 'seismic', 'vim':-60, 'vmax':-20, 'name': 'LDR'}, \n",
    "               'ldrhh14SP': {'cmp': 'seismic', 'vim':-60, 'vmax':-20,  'name': 'LDR'}}\n",
    "        \n",
    "        x = ds.range * np.sin(np.deg2rad(ds.azimuth)) * ds.DR.values[0]\n",
    "        y = ds.alt3D \n",
    "        \n",
    "        try:\n",
    "        \n",
    "            ax1 = axs[0][0].pcolormesh(x[:, :], y[:, :], ds.zhh14[:, :], cmap=var['zhh14']['cmp'], \n",
    "                              vmin=var['zhh14']['vim'], vmax=var['zhh14']['vmax'], shading='auto')\n",
    "\n",
    "            ax2 = axs[0][1].pcolormesh(x[:, :], y[:, :], ds.zhh14SP[:, :], cmap=var['zhh14SP']['cmp'], \n",
    "                              vmin=var['zhh14SP']['vim'], vmax=var['zhh14SP']['vmax'], shading='auto')\n",
    "\n",
    "            ax3 = axs[1][0] .pcolormesh(x[:, :], y[:, :], ds.vel14[:, :], cmap=var['vel14']['cmp'], \n",
    "                              vmin=var['vel14']['vim'], vmax=var['vel14']['vmax'], shading='flat', rasterized=True)\n",
    "\n",
    "            ax4 = axs[1][1].pcolormesh(x[:, :], y[:, :], ds.vel14SP[:, :], cmap=var['vel14SP']['cmp'], \n",
    "                              vmin=var['vel14SP']['vim'], vmax=var['vel14SP']['vmax'], shading='flat', rasterized=True)\n",
    "\n",
    "            ax5 = axs[2][0].pcolormesh(x[:, :], y[:, :], ds.ldrhh14[:, :], cmap=var['ldrhh14']['cmp'], \n",
    "                              vmin=var['ldrhh14']['vim'], vmax=var['ldrhh14']['vmax'], shading='flat', rasterized=True)\n",
    "\n",
    "            ax6 = axs[2][1].pcolormesh(x[:, :], y[:, :], ds.ldrhh14SP[:, :], cmap=var['ldrhh14SP']['cmp'], \n",
    "                              vmin=var['ldrhh14SP']['vim'], vmax=var['ldrhh14SP']['vmax'], shading='flat', rasterized=True)\n",
    "\n",
    "            plt.colorbar(ax1, ax=[axs[0][0], axs[0][1]], pad=0.05, label=f\"{var['zhh14']['name']} ({ds['zhh14'].units})\")\n",
    "            plt.colorbar(ax3, ax=[axs[1][0], axs[1][1]], pad=0.05, label=f\"{var['vel14']['name']} ({ds['vel14'].units})\")\n",
    "            plt.colorbar(ax5, ax=[axs[2][0], axs[2][1]], pad=0.05, label=f\"{var['ldrhh14']['name']} ({ds['ldrhh14'].units})\")\n",
    "            save = f\"{np.datetime_as_string(ds.time.values, unit='s').replace(':','-')}\"\n",
    "            t = ds.time.values\n",
    "        except ValueError:\n",
    "            \n",
    "            ax1 = axs[0][0].pcolormesh(x[:, :, 0], y[:, :, 0], ds.zhh14[:, :, 0], cmap=var['zhh14']['cmp'], \n",
    "                              vmin=var['zhh14']['vim'], vmax=var['zhh14']['vmax'], shading='auto')\n",
    "\n",
    "            ax2 = axs[0][1].pcolormesh(x[:, :, 0], y[:, :, 0], ds.zhh14SP[:, :, 0], cmap=var['zhh14SP']['cmp'], \n",
    "                              vmin=var['zhh14SP']['vim'], vmax=var['zhh14SP']['vmax'], shading='nearest')\n",
    "\n",
    "            ax3 = axs[1][0] .pcolormesh(x[:, :, 0], y[:, :, 0], ds.vel14[:, :, 0], cmap=var['vel14']['cmp'], \n",
    "                              vmin=var['vel14']['vim'], vmax=var['vel14']['vmax'], shading='nearest', rasterized=True)\n",
    "\n",
    "            ax4 = axs[1][1].pcolormesh(x[:, :, 0], y[:, :, 0], ds.vel14SP[:, :, 0], cmap=var['vel14SP']['cmp'], \n",
    "                              vmin=var['vel14SP']['vim'], vmax=var['vel14SP']['vmax'], shading='nearest', rasterized=True)\n",
    "\n",
    "            ax5 = axs[2][0].pcolormesh(x[:, :, 0], y[:, :, 0], ds.ldrhh14[:, :, 0], cmap=var['ldrhh14']['cmp'], \n",
    "                              vmin=var['ldrhh14']['vim'], vmax=var['ldrhh14']['vmax'], shading='nearest', rasterized=True)\n",
    "\n",
    "            ax6 = axs[2][1].pcolormesh(x[:, :, 0], y[:, :, 0], ds.ldrhh14SP[:, :, 0], cmap=var['ldrhh14SP']['cmp'], \n",
    "                              vmin=var['ldrhh14SP']['vim'], vmax=var['ldrhh14SP']['vmax'], shading='nearest', rasterized=True)\n",
    "\n",
    "            plt.colorbar(ax1, ax=[axs[0][0], axs[0][1]], pad=0.05, label=f\"{var['zhh14']['name']} ({ds['zhh14'].units})\")\n",
    "            plt.colorbar(ax3, ax=[axs[1][0], axs[1][1]], pad=0.05, label=f\"{var['vel14']['name']} ({ds['vel14'].units})\")\n",
    "            plt.colorbar(ax5, ax=[axs[2][0], axs[2][1]], pad=0.05, label=f\"{var['ldrhh14']['name']} ({ds['ldrhh14'].units})\")\n",
    "            save = f\"{np.datetime_as_string(ds.time.values[0], unit='s').replace(':','-')}\"\n",
    "            t = ds.time\n",
    "#         for i, ax  in enumerate(axs.flat):\n",
    "#             key = list(var.keys())\n",
    "#             x = ds.range * np.sin(np.deg2rad(ds.azimuth)) * ds.DR.values[0]\n",
    "#             y = ds.alt3D \n",
    "#             data =  ds[key[i]]\n",
    "#             data = data.where(data != -9999.)\n",
    "#             im = ax.pcolormesh(x[:, :, 0], y[:, :, 0], data[:, :, 0], cmap=var[key[i]]['cmp'], \n",
    "#                           vmin=var[key[i]]['vim'], vmax=var[key[i]]['vmax'], shading='auto')\n",
    "#             ax.set_ylabel('Height (m)')\n",
    "#             ax.set_xlabel('Cross track (m)')\n",
    "#             ax.set_title(\"\\n\".join(wrap(ds[key[i]].notes, 30)))\n",
    "#             ax.set_ylim(0, 8000)\n",
    "#             if i % 2 != 0:\n",
    "#                 plt.colorbar(im, ax=[axs.flat[i-1:i+1]], pad=0.05, \n",
    "#                              label=f\"{var[key[i]]['name']} ({ds[key[i]].units})\")\n",
    "\n",
    "        title = f\"{np.datetime_as_string(ds.time.values, unit='s')}\\n\"\n",
    "        fig.suptitle(title, fontsize=16, y=.92)\n",
    "        fig.savefig(f\"{path_proj}/results/plots/test/{save}.jpg\", format='jpg')\n",
    "        plt.close('all')\n",
    "        del ds\n",
    "        return t\n",
    "    \n",
    "# @dask.delayed\n",
    "def load(_time):\n",
    "    ds_xr = xr.open_zarr(f'{path_data}/zarr_rckd/KUsKAs_Wn/lores.zarr')\n",
    "    ds_xr = ds_xr.sel(time=~ds_xr.get_index(\"time\").duplicated())\n",
    "    var2load = ['range','alt3D', 'azimuth', 'zhh14', 'zhh14SP', 'vel14', \n",
    "               'vel14SP', 'ldrhh14', 'ldrhh14SP', 'DR']\n",
    "    w = ds_xr[var2load].sel(time=_time)\n",
    "    return w\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "!scancel -u alfonso8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = SLURMCluster(queue=\"seseml\",\n",
    "                       memory='100GB',\n",
    "                       cores=40,\n",
    "                       processes=10,\n",
    "                       walltime='24:40:00',\n",
    "                       scheduler_options={'host': '172.22.179.3:7222', \n",
    "                                          'dashboard_address': '7999',\n",
    "                                          },                        \n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dbbd31b366c14edf88686e7bcef78f84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tab(children=(HTML(value='<div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-outpu…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cluster.scale(jobs=8)\n",
    "# cluster.adapt(maximum_jobs=8)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)\n",
      "            548708    seseml dask-wor alfonso8 PD       0:00      1 (None)\n",
      "            548707    seseml dask-wor alfonso8 PD       0:00      1 (None)\n",
      "            548706    seseml dask-wor alfonso8 PD       0:00      1 (None)\n",
      "            548705    seseml dask-wor alfonso8 PD       0:00      1 (None)\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "squeue -u alfonso8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-6341ef73-786a-11ec-8e86-b083fed70206</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.SLURMCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://172.22.179.3:7999/status\" target=\"_blank\">http://172.22.179.3:7999/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">SLURMCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">b555944c</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://172.22.179.3:7999/status\" target=\"_blank\">http://172.22.179.3:7999/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-04877f8c-12aa-4656-9326-38d34bed8c9f</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://172.22.179.3:7222\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://172.22.179.3:7999/status\" target=\"_blank\">http://172.22.179.3:7999/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://172.22.179.3:7222' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11366\n"
     ]
    }
   ],
   "source": [
    "ds_prop = pd.read_csv(f'{path_proj}/results/all_px_202111070706.csv')\n",
    "# ds_prop.dates = pd.to_datetime(ds_prop.dates)\n",
    "ds_prop.rename(columns={'Unnamed: 0': 'dates'}, inplace=True)\n",
    "times = list(ds_prop.dates)\n",
    "print(len(times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "78800"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.28 s, sys: 128 ms, total: 1.41 s\n",
      "Wall time: 1.32 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ds_xr = xr.open_zarr(f'{path_data}/zarr_rckd/KUsKAs_Wn/lores.zarr')\n",
    "# ds_xr = xr.open_zarr(f'{path_data}/zarr/KUsKAs_Wn/lores.zarr')\n",
    "ds_xr = ds_xr.sel(time=~ds_xr.get_index(\"time\").duplicated())\n",
    "len(ds_xr.time.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_sel = slice('2019-09-15 03:11:00','2019-09-17 03:50:00')\n",
    "# time_sel = slice('2019-09-16 03:11:00','2019-09-16 10:50:00')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11059"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds_xr = ds_xr.sel(time=time_sel)\n",
    "len(ds_xr.time.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "var2load = ['range','alt3D', 'azimuth', 'zhh14', 'zhh14SP', 'vel14', \n",
    "               'vel14SP', 'ldrhh14', 'ldrhh14SP', 'DR']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = ds_xr[var2load].chunk({'time':1, 'cross_track':-1, 'range':-1})\\\n",
    "        .map_blocks(multiple_plot, template=xr.ones_like(ds_xr.time).chunk({\"time\": 1}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4min 32s, sys: 22.8 s, total: 4min 55s\n",
      "Wall time: 4min 54s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "res = tasks.compute()\n",
    "## entra en el Except"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_delayed = [delayed(multiple_plot)(ds_xr[var2load].sel(time=i)) for i in ds_xr.time]\n",
    "len(task_delayed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "res = dask.compute(*task_delayed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:camp2ex_proj] *",
   "language": "python",
   "name": "conda-env-camp2ex_proj-py"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}