#!/usr/bin/env python
# coding: utf-8

# # How to process CfRadial1 Volume
# 
# along the example of cluttermap creation

# See question on [openradar discourse](https://openradar.discourse.group/t/how-to-filter-rays/).

# We will make use of the following Open Source Software Packages:
# 
# - [xarray](https://docs.xarray.dev)
# - [xarray-datatree](https://xarray-datatree.readthedocs.io)
# - [xradar](https://docs.openradarscience.org/projects/xradar/)
# - [wradlib](https://docs.wradlib.org/)
# - [matplotlib](https://matplotlib.org/)

# In[1]:


import xarray as xr
import datatree
import xradar
import wradlib as wrl
import matplotlib.pyplot as plt


# In[2]:


fname = "BHP210601110223.nc"


# ## Define wrapper function for clutter extraction

# In[3]:


def extract_clutter(da, refl, cmap, wsize=3, thrsnorain=0, tr1=6.0, n_p=6, tr2=1.3, rm_nans=False):
    if refl in da.variables:
        da = da.assign({cmap: xr.apply_ufunc(
            wrl.clutter.filter_gabella,
            da[refl],
            input_core_dims=[["azimuth", "range"]],
            output_core_dims=[["azimuth", "range"]],
            dask="parallelized",
            kwargs=dict(
                wsize=wsize,
                thrsnorain=thrsnorain,
                tr1=tr1,
                n_p=n_p,
                tr2=tr2,
                rm_nans=rm_nans,
            ),
        )})
    return da


# ## Open CfRadial1 data file

# We make use of our new openradar community package ``xradar`` which builds on ``xarray`` and ``xarray-datatree``.

# In[4]:


vol = xradar.io.open_cfradial1_datatree(fname)


# In[5]:


display(vol)


# ## Map clutter extraction over volume
# 
# This uses `xarray-datatree` functionality as well as `xarray` and `wradlib`.

# In[6]:


vol = vol.map_over_subtree(extract_clutter, "reflectivity", "CMAP", wsize=3, thrsnorain=0.0, tr1=14, n_p=2, tr2=1.5, rm_nans=False)


# ## Export to new Cfradial2 Volume
# 
# Uses `xradar`.

# In[7]:


new_name = "new_cfradial2_volume.nc"
xradar.io.export.to_cfradial2(vol, new_name)


# ## Reopen Cfradial2 Volume
# 
# Uses `xarray-datatree`.

# In[8]:


nvol = datatree.open_datatree(new_name)
display(nvol)


# ## Georeference Volume
# 
# We use `xradar`-Accessor here. The sweeps are georeferenced with azimuthal equidistant projection.

# In[9]:


nvol = nvol.xradar.georeference()


# ## Plotting
# 
# Uses `xarray` provided plotting using `matplotlib`.

# ### Get first sweep Dataset by name

# In[10]:


swpx = nvol["sweep_0"].ds
display(swpx)


# ### Plot first sweep

# In[11]:


fig = plt.figure(figsize=(15, 12))

ax1 = fig.add_subplot(221)
swpx.reflectivity.plot(x="x", y="y", ax=ax1, vmin=0, vmax=60)
ax1.set_title("Reflectivity raw")

ax2 = fig.add_subplot(222)
swpx.CMAP.plot(x="x", y="y", ax=ax2)
ax2.set_title("Cluttermap")

ax3 = fig.add_subplot(223)
swpx.reflectivity.where(swpx.CMAP == 1).plot(
    x="x", y="y", ax=ax3, vmin=0, vmax=60
)
ax3.set_title("Clutter")

ax4 = fig.add_subplot(224)
swpx.reflectivity.where(swpx.CMAP < 1).plot(
    x="x", y="y", ax=ax4, vmin=0, vmax=60
)
ax4.set_title("Reflectivity clutter removed")


# ### Get second sweep Dataset by name

# In[12]:


swpx = nvol["sweep_1"].ds
display(swpx)


# ### PPlot Second Sweep

# In[13]:


fig = plt.figure(figsize=(15, 12))

ax1 = fig.add_subplot(221)
swpx.reflectivity.plot(x="x", y="y", ax=ax1, vmin=0, vmax=60)
ax1.set_title("Reflectivity raw")

ax2 = fig.add_subplot(222)
swpx.CMAP.plot(x="x", y="y", ax=ax2)
ax2.set_title("Cluttermap")

ax3 = fig.add_subplot(223)
swpx.reflectivity.where(swpx.CMAP == 1).plot(
    x="x", y="y", ax=ax3, vmin=0, vmax=60
)
ax3.set_title("Clutter")

ax4 = fig.add_subplot(224)
swpx.reflectivity.where(swpx.CMAP < 1).plot(
    x="x", y="y", ax=ax4, vmin=0, vmax=60
)
ax4.set_title("Reflectivity clutter removed")