{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Wradlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import radarx as rx\n",
    "import xarray as xr\n",
    "import xradar as xd\n",
    "\n",
    "# from osgeo import osr\n",
    "import wradlib as wrl\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Define Funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_radar(ds):\n",
    "    \"\"\"Filters radar data for valid reflectivity values.\"\"\"\n",
    "    return ds.where((ds.reflectivity >= 0) & (ds.reflectivity <= 70))\n",
    "\n",
    "\n",
    "def setup_grid(dtree, horiz_res=2000, vert_res=500, maxrange=275e3, maxalt=15e3):\n",
    "    \"\"\"Creates a 3D grid for wradlib.\"\"\"\n",
    "    # sitecoords = (\n",
    "    #     dtree[\"sweep_0\"].longitude.values,\n",
    "    #     dtree[\"sweep_0\"].latitude.values,\n",
    "    #     dtree[\"sweep_0\"].altitude.values,\n",
    "    # )\n",
    "    #\n",
    "    #     proj = osr.SpatialReference()\n",
    "    #     proj.ImportFromEPSG(4326)  # WGS 84 (EPSG:4326)\n",
    "    #     trgxyz, trgshape = wrl.vpr.make_3d_grid(\n",
    "    #         sitecoords, proj, maxrange, maxalt, horiz_res, vert_res)\n",
    "    # ------------------\n",
    "    #     This seems to crash RTD build, so trying something else\n",
    "    #     see log https://app.readthedocs.org/api/v2/build/26513933.txt\n",
    "    # -----------------\n",
    "\n",
    "    x_lims = (-300e3, 250e3)\n",
    "    y_lims = (-250e3, 300e3)\n",
    "\n",
    "    lat, lon, x, y, z, trg_crs = rx.grid.make_3d_grid(\n",
    "        dtree[\"sweep_0\"].to_dataset(),\n",
    "        x_lim=x_lims,\n",
    "        y_lim=y_lims,\n",
    "        x_step=horiz_res,\n",
    "        y_step=horiz_res,\n",
    "        z_lim=(0, maxalt),\n",
    "        z_step=vert_res,\n",
    "    )\n",
    "\n",
    "    trgshape = len(z), len(y), len(x)\n",
    "    trgxyz = wrl.util.gridaspoints(z, y, x)\n",
    "    return trgxyz, trgshape\n",
    "\n",
    "\n",
    "def process_sweeps(raw_dt, data_var=\"reflectivity\"):\n",
    "    \"\"\"Processes sweep data and prepares for gridding.\"\"\"\n",
    "    swp_list, data_list = [], []\n",
    "\n",
    "    for swp in raw_dt.match(\"sweep_*\"):\n",
    "        ds = raw_dt[swp].to_dataset()\n",
    "        xyz = (\n",
    "            xr.concat(\n",
    "                [\n",
    "                    ds.coords[\"x\"].reset_coords(drop=True),\n",
    "                    ds.coords[\"y\"].reset_coords(drop=True),\n",
    "                    ds.coords[\"z\"].reset_coords(drop=True),\n",
    "                ],\n",
    "                \"xyz\",\n",
    "            )\n",
    "            .stack(npoints=(\"azimuth\", \"range\"))\n",
    "            .transpose(..., \"xyz\")\n",
    "        )\n",
    "        swp_list.append(xyz)\n",
    "        data = ds[data_var].stack(npoints=(\"azimuth\", \"range\"))\n",
    "        data_list.append(data)\n",
    "\n",
    "    xyz = xr.concat(swp_list, \"npoints\")\n",
    "    data = xr.concat(data_list, \"npoints\")\n",
    "    return xyz, data\n",
    "\n",
    "\n",
    "def interpolate_to_grid(xyz, trgxyz, data, trgshape, maxrange, minelev, maxelev):\n",
    "    \"\"\"Interpolates radar data to a Cartesian 3D volume grid.\"\"\"\n",
    "    gridder = wrl.vpr.CAPPI(\n",
    "        xyz.values,\n",
    "        trgxyz,\n",
    "        maxrange=maxrange,\n",
    "        minelev=minelev,\n",
    "        maxelev=maxelev,\n",
    "    )\n",
    "    vol = np.ma.masked_invalid(gridder(data.values).reshape(trgshape))\n",
    "    return vol\n",
    "\n",
    "\n",
    "def create_dataset(\n",
    "    vol, trgxyz, trgshape, data, proj_crs, dtree, data_var=\"reflectivity\"\n",
    "):\n",
    "    \"\"\"Creates an Xarray dataset from gridded radar data.\"\"\"\n",
    "    trgx, trgy, trgz = (\n",
    "        trgxyz[:, 0].reshape(trgshape)[0, 0, :],\n",
    "        trgxyz[:, 1].reshape(trgshape)[0, :, 0],\n",
    "        trgxyz[:, 2].reshape(trgshape)[:, 0, 0],\n",
    "    )\n",
    "\n",
    "    lon, lat = rx.utils.cartesian_to_geographic_aeqd(\n",
    "        trgx,\n",
    "        trgy,\n",
    "        data.longitude.values,\n",
    "        data.latitude.values,\n",
    "        xd.georeference.get_earth_radius(proj_crs, data.latitude.values),\n",
    "    )\n",
    "\n",
    "    ds_wrl = xr.DataArray(\n",
    "        data=vol,\n",
    "        coords={\"z\": trgz, \"y\": trgy, \"x\": trgx},\n",
    "        dims=(\"z\", \"y\", \"x\"),\n",
    "        name=data_var,\n",
    "    ).to_dataset()\n",
    "\n",
    "    ds_wrl[\"time\"] = data.time.mean()\n",
    "    ds_wrl.attrs = dtree.attrs\n",
    "    ds_wrl[\"latitude\"] = data[\"latitude\"]\n",
    "    ds_wrl[\"longitude\"] = data[\"longitude\"]\n",
    "    ds_wrl[\"lon\"] = xr.DataArray(lon, dims=[\"x\"])\n",
    "    ds_wrl[\"lat\"] = xr.DataArray(lat, dims=[\"y\"])\n",
    "    ds_wrl = ds_wrl.set_coords([\"lon\", \"lat\"])\n",
    "    return ds_wrl"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "# Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    # Parameters\n",
    "    filename = \"KSGF20180612_083109_V06.nc\"\n",
    "\n",
    "    # Grid Setup\n",
    "    h_res, v_res = 2000, 500\n",
    "    maxrange, maxalt = 275e3, 15e3\n",
    "    minelev, maxelev = 0.2, 21.0\n",
    "\n",
    "    data_var = \"reflectivity\"\n",
    "\n",
    "    # Read and process data\n",
    "    dtree = xd.io.open_cfradial1_datatree(filename)\n",
    "    dtree = rx.utils.combine_nexrad_sweeps(dtree)\n",
    "    dtree = dtree.xradar.map_over_sweeps(filter_radar)\n",
    "    dtree = dtree.xradar.georeference()\n",
    "\n",
    "    tstart = datetime.now()\n",
    "\n",
    "    trgxyz, trgshape = setup_grid(\n",
    "        dtree, horiz_res=h_res, vert_res=v_res, maxrange=maxrange, maxalt=maxalt\n",
    "    )\n",
    "\n",
    "    raw_dt = dtree.xradar.map_over_sweeps(rx.utils.get_geocoords)\n",
    "    proj_crs = xd.georeference.get_crs(raw_dt[\"sweep_0\"].ds)\n",
    "\n",
    "    xyz, data = process_sweeps(raw_dt, data_var=data_var)\n",
    "\n",
    "    vol = interpolate_to_grid(xyz, trgxyz, data, trgshape, maxrange, minelev, maxelev)\n",
    "    ds_wrl = create_dataset(\n",
    "        vol, trgxyz, trgshape, data, proj_crs, dtree, data_var=data_var\n",
    "    )\n",
    "\n",
    "    print(\"Wradlib gridding took:\", datetime.now() - tstart)\n",
    "    display(ds_wrl)\n",
    "\n",
    "    # Diagnostic plot\n",
    "    ds_wrl.radarx.plot_max_cappi(\n",
    "        data_var=data_var,\n",
    "        vmin=-10,\n",
    "        vmax=70,\n",
    "        range_rings=True,\n",
    "        add_map=True,\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}