{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CfRadial1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import cmweather\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xradar as xd\n",
    "import datatree as xt\n",
    "from open_radar_data import DATASETS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download\n",
    "\n",
    "Fetching CfRadial1 radar data file from [open-radar-data](https://github.com/openradar/open-radar-data) repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = DATASETS.fetch(\"cfrad.20080604_002217_000_SPOL_v36_SUR.nc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar = xd.io.open_cfradial1_datatree(filename, first_dim=\"auto\")\n",
    "display(radar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Azimuth vs. Range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar.sweep_0.DBZ.plot(cmap=\"ChaseSpectral\", vmin=-10, vmax=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Time vs. Range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar.sweep_0.DBZ.swap_dims({\"azimuth\": \"time\"}).sortby(\"time\").plot(\n",
    "    cmap=\"ChaseSpectral\", vmin=-10, vmax=70\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Georeference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar = radar.xradar.georeference()\n",
    "display(radar)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot PPI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar[\"sweep_0\"][\"DBZ\"].plot(x=\"x\", y=\"y\", cmap=\"ChaseSpectral\", vmin=-10, vmax=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter\n",
    "\n",
    "Apply basic reflectivity filter. This is just a demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ref_filter(dtree, sweep=\"sweep_0\", field=\"DBZ\"):\n",
    "    ds = dtree[sweep].where((dtree[sweep][field] >= -10) & (dtree[sweep][field] <= 70))\n",
    "    red_patch = ds.where(\n",
    "        (\n",
    "            (ds[field] >= ds[field].max().values - 0.5)\n",
    "            & (ds[field] <= ds[field].max().values + 0.5)\n",
    "        ),\n",
    "        drop=True,\n",
    "    )\n",
    "    rmin, rmax = int(red_patch.range.min().values - 150), int(\n",
    "        red_patch.range.max().values + 150\n",
    "    )\n",
    "    out_of_range_mask = (ds.range < rmin) | (ds.range > rmax)\n",
    "    ds[field] = ds[field].where(out_of_range_mask)\n",
    "    # Interpolate missing values using the slinear method along the 'range' dimension\n",
    "    ds[field] = ds[field].interpolate_na(dim=\"range\", method=\"slinear\")\n",
    "    dtree[sweep][f\"corr_{field}\"] = ds[field].copy()\n",
    "    return dtree[sweep]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swp0 = ref_filter(radar, sweep=\"sweep_0\", field=\"DBZ\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "swp0.corr_DBZ.plot(x=\"x\", y=\"y\", cmap=\"ChaseSpectral\", vmin=-10, vmax=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter full volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize an empty DataTree\n",
    "result_tree = xt.DataTree()\n",
    "\n",
    "for sweep in radar.sweep_group_name.values:\n",
    "    corrected_data = ref_filter(radar, sweep, field=\"DBZ\")\n",
    "\n",
    "    # Convert the xarray Dataset to a DataTree and add it to the result_tree\n",
    "    data_tree = xt.DataTree.from_dict(corrected_data.to_dict())\n",
    "\n",
    "    # Copy the contents of data_tree into result_tree\n",
    "    for key, value in data_tree.items():\n",
    "        result_tree[key] = value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radar.sweep_6.corr_DBZ.plot(x=\"x\", y=\"y\", cmap=\"ChaseSpectral\", vmin=-10, vmax=70)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export\n",
    "\n",
    "Export to CfRadial1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xd.io.to_cfradial1(dtree=radar, filename=\"cfradial1_qced.nc\", calibs=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?xd.io.to_cfradial1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Note \n",
    "\n",
    "If `filename` is `None` in the `xd.io.to_cfradial1` function, it will automatically generate a<br>\n",
    "filename using the instrument name and the first available timestamp from the data.\n"
   ]
  }
 ],
 "metadata": {
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}