"""
=======================================
Convective-Stratiform classification
=======================================
This example shows how to use the updated convective stratiform classifcation algorithm. We show 3 examples,
a summer convective example, an example from Hurricane Ian, and an example from a winter storm.

References
----------
Steiner, M. R., R. A. Houze Jr., and S. E. Yuter, 1995: Climatological
Characterization of Three-Dimensional Storm Structure from Operational
Radar and Rain Gauge Data. J. Appl. Meteor., 34, 1978-2007.
https://doi.org/10.1175/1520-0450(1995)034<1978:CCOTDS>2.0.CO;2.

Yuter, S. E., and R. A. Houze, Jr., 1997: Measurements of raindrop size
distributions over the Pacific warm pool and implications for Z-R relations.
J. Appl. Meteor., 36, 847-867.
https://doi.org/10.1175/1520-0450(1997)036%3C0847:MORSDO%3E2.0.CO;2

Yuter, S. E., R. A. Houze, Jr., E. A. Smith, T. T. Wilheit, and E. Zipser,
2005: Physical characterization of tropical oceanic convection observed in
KWAJEX. J. Appl. Meteor., 44, 385-415. https://doi.org/10.1175/JAM2206.1

Rasmussen, R., M. Dixon, S. Vasiloff, F. Hage, S. Knight, J. Vivekanandan,
and M. Xu, 2003: Snow Nowcasting Using a Real-Time Correlation of Radar
Reflectivity with Snow Gauge Accumulation. J. Appl. Meteorol. Climatol., 42, 20–36.
https://doi.org/10.1175/1520-0450(2003)042%3C0020:SNUART%3E2.0.CO;2
"""

print(__doc__)

# Author: Laura Tomkins (lmtomkin@ncsu.edu)
# License: BSD 3 clause


import pyart
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import cartopy.crs as ccrs

######################################
# **How the algorithm works**
# This first section describes how the convective-stratiform algorithm works (see references for full details). This
# algorithm is a feature detection algorithm and classifies fields as "convective" or "stratiform". The algorithm is
# designed to detect features in a reflectivity field but can also detect features in fields such as rain rate or
# snow rate. In this section we describe the steps of the convective stratiform algorithm and the variables used in
# the function.
# The first step of the algorithm calculates a background average of the field with a circular footprint using a radius
# provided by `bkg_rad_km`. A larger radius will yield a smoother field. The radius needs to be at least double the
# grid spacing, but we recommend at least three times the grid spacing. If using reflectivity, `dB_averaging` should be set
# to True to convert reflectivity to linear Z before averaging, False for rescaled fields such as rain or snow rate.
# `calc_thres` determines the minimum fraction of a circle that is considered in the background average calculation
# (default is 0.75, so the points along the edges where there is less than 75% of a full circle of data,
# the algorithm is not run).
# Once the background average has been calculated, the original field is compared to the background average.  In
# order for points to be considered "convective cores" they must exceed the background value by a certain value or
# simply be greater than the `always_core_thres`. This value is determined by either a cosine scheme, or a scalar
# value (i.e. the reflectivity value must be X times the background value (multiplier; `use_addition=False`),
# or  X greater than the background value (`use_addition=True`) where X is the `scalar_diff`).
# `use_cosine` determines if a cosine scheme or scalar scheme is to be used. If `use_cosine` is True,
# then the `max_diff` and `zero_diff_cos_val` come into use. These values define the cosine scheme that is used  to
# determine the minimum difference between the background and reflectivity value in order for a core to be
# identified. `max_diff` is the maximum difference between the field and the background for a core to be identified,
# or where the cosine function crosses the y-axis. The `zero_diff_cos_val` is where the difference between the  field
# and the background is zero, or where the cosine function crosses the x-axis. Note, if
# `always_core_thres` < `zero_diff_cos_val`, `zero_diff_cos_val` only helps define the shape of the cosine curve and
# all values greater than `always_core_thres` will be considered a convective core. If
# `always_core_thres` > `zero_diff_cos_val` then all values greater than `zero_diff_cos_val` will be considered a
# convective core. We plot some examples of the schemes below:

# plot cosine scheme
pyart.graph.plot_convstrat_scheme(always_core_thres=30, use_cosine=True, max_diff=5, zero_diff_cos_val=45)
# when zero_diff_cos_val is greater than always_core_thres, the difference becomes zero at the zero_diff_cos_val
pyart.graph.plot_convstrat_scheme(always_core_thres=55, use_cosine=True, max_diff=5, zero_diff_cos_val=45)
# alternatively, we can use a simpler scalar difference instead of a cosine scheme
pyart.graph.plot_convstrat_scheme(always_core_thres=40, use_cosine=False, max_diff=None,
                                  zero_diff_cos_val=None, use_addition=True, scalar_diff=2)
# if you are interested in picking up weak features, you can also use the scalar difference as a multiplier instead,
# so very weak features do not have to be that different from the background to be classified as convective.
pyart.graph.plot_convstrat_scheme(always_core_thres=40, use_cosine=False, max_diff=None,
                                  zero_diff_cos_val=None, use_addition=False, scalar_diff=2)

######################################
# Once the cores are identified, there is an option to remove speckles (`remove_small_objects`) smaller than a  given
# size (`min_km2_size`).
# After the convective cores are identified, We then incorporate convective radii using
# `val_for_max_conv_rad` and `max_conv_rad_km`. The convective radii act as a dilation and are used to classify
# additional points around the cores as convective that may not have been identified previously.  The
# `val_for_max_conv_rad` is the value where the maximum convective radius is applied and the `max_conv_rad_km` is the
# maximum convective radius. Values less than the `val_for_max_conv_rad` are assigned a convective radius using a step
# function.
# Finally, the points are classified as NOSFCECHO (threshold set with `min_dBZ_used`; 0), WEAKECHO (threshold set with
# `weak_echo_thres`; 3), SF (stratiform; 1), CONV (convective; 2).

######################################
# **Classification of summer convective example**
# Our first example classifies echo from a summer convective event. We use a cosine scheme to classify the convective
# points.

# Now let's do a classification with our parameters
# read in file
filename = pyart.testing.get_test_data('swx_20120520_0641.nc')
radar = pyart.io.read(filename)

# extract the lowest sweep
radar = radar.extract_sweeps([0])

# interpolate to grid
grid = pyart.map.grid_from_radars(
    (radar,), grid_shape=(1, 201, 201),
    grid_limits=((0, 10000), (-50000.0, 50000.0), (-50000.0, 50000.0)),
    fields=['reflectivity_horizontal'])

# get dx dy
dx = grid.x['data'][1] - grid.x['data'][0]
dy = grid.y['data'][1] - grid.y['data'][0]

# convective stratiform classification
convsf_dict = pyart.retrieve.conv_strat_yuter(grid, dx, dy, refl_field='reflectivity_horizontal', always_core_thres=40,
                                              bkg_rad_km=20, use_cosine=True, max_diff=5, zero_diff_cos_val=55,
                                              weak_echo_thres=10, max_conv_rad_km=2)

# add to grid object
# mask zero values (no surface echo)
convsf_masked = np.ma.masked_equal(convsf_dict['convsf']['data'], 0)
# mask 3 values (weak echo)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
# add dimension to array to add to grid object
convsf_dict['convsf']['data'] = convsf_masked[None,:,:]
# add field
grid.add_field('convsf', convsf_dict['convsf'], replace_existing=True)

# create plot using GridMapDisplay
# plot variables
display = pyart.graph.GridMapDisplay(grid)
magma_r_cmap = plt.get_cmap('magma_r')
ref_cmap = mcolors.LinearSegmentedColormap.from_list('ref_cmap', magma_r_cmap(np.linspace(0, 0.9, magma_r_cmap.N)))
projection = ccrs.LambertConformal(central_latitude=radar.latitude['data'][0],
                                   central_longitude=radar.longitude['data'][0])
######################################
# Note how the convective stratiform field has less data around the edges compared to the reflectivity
# field. This is because the function is designed to only compute the background radius where the footprint has 75%
# of data, so along the edges there is not enough data to fulfill this requirement. The footprint percentage can be
# changed using the variable `calc_thres`.

# plot
plt.figure(figsize=(10,4))
ax1=plt.subplot(1,2,1, projection=projection)
display.plot_grid('reflectivity_horizontal', vmin=5, vmax=45, cmap=ref_cmap, projection=projection,
                  transform=ccrs.PlateCarree(), ax=ax1)
ax2=plt.subplot(1,2,2, projection=projection)
display.plot_grid('convsf', vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3), projection=projection, ax=ax2,
                  transform=ccrs.PlateCarree(), ticks=[1/3, 1, 5/3], ticklabs=['', 'Stratiform', 'Convective'])
plt.show()


######################################
# In addition to the default convective-stratiform classification, the function also returns an underestimate
# (convsf_under) and an overestimate (convsf_over) to take into consideration the uncertainty when choosing
# classification parameters. The under and overestimate use the same parameters, but vary the input field by a
# certain value (default is 5 dBZ, can be changed with `estimate_offset`). The estimation can be turned off (
# `estimate_flag=False`), but we recommend keeping it turned on.

# mask weak echo and no surface echo
convsf_masked = np.ma.masked_equal(convsf_dict['convsf']['data'], 0)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
convsf_dict['convsf']['data'] = convsf_masked
# underest.
convsf_masked = np.ma.masked_equal(convsf_dict['convsf_under']['data'], 0)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
convsf_dict['convsf_under']['data'] = convsf_masked
# overest.
convsf_masked = np.ma.masked_equal(convsf_dict['convsf_over']['data'], 0)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
convsf_dict['convsf_over']['data'] = convsf_masked

# Plot each estimation
plt.figure(figsize=(10,4))
ax1=plt.subplot(131)
ax1.pcolormesh(convsf_dict['convsf']['data'][0,:,:], vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3))
ax1.set_title('Best estimate')
ax1.set_aspect('equal')
ax2=plt.subplot(132)
ax2.pcolormesh(convsf_dict['convsf_under']['data'], vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3))
ax2.set_title('Underestimate')
ax2.set_aspect('equal')
ax3=plt.subplot(133)
ax3.pcolormesh(convsf_dict['convsf_over']['data'], vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3))
ax3.set_title('Overestimate')
ax3.set_aspect('equal')
plt.show()

######################################
# **Tropical example**
# Let's get a NEXRAD file from Hurricane Ian

# Read in file
nexrad_file = 's3://noaa-nexrad-level2/2022/09/28/KTBW/KTBW20220928_190142_V06'
radar = pyart.io.read_nexrad_archive(nexrad_file)

# extract the lowest sweep
radar = radar.extract_sweeps([0])

# interpolate to grid
grid = pyart.map.grid_from_radars(
    (radar,), grid_shape=(1, 201, 201),
    grid_limits=((0, 10000), (-200000.0, 200000.0), (-200000.0, 200000.0)),
    fields=['reflectivity'])

# get dx dy
dx = grid.x['data'][1] - grid.x['data'][0]
dy = grid.y['data'][1] - grid.y['data'][0]

# convective stratiform classification
convsf_dict = pyart.retrieve.conv_strat_yuter(grid, dx, dy, refl_field='reflectivity', always_core_thres=40,
                                              bkg_rad_km=20, use_cosine=True, max_diff=3, zero_diff_cos_val=55,
                                              weak_echo_thres=5, max_conv_rad_km=2, estimate_flag=False)

# add to grid object
# mask zero values (no surface echo)
convsf_masked = np.ma.masked_equal(convsf_dict['convsf']['data'], 0)
# mask 3 values (weak echo)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
# add dimension to array to add to grid object
convsf_dict['convsf']['data'] = convsf_masked[None,:,:]
# add field
grid.add_field('convsf', convsf_dict['convsf'], replace_existing=True)

# create plot using GridMapDisplay
# plot variables
display = pyart.graph.GridMapDisplay(grid)
magma_r_cmap = plt.get_cmap('magma_r')
ref_cmap = mcolors.LinearSegmentedColormap.from_list('ref_cmap', magma_r_cmap(np.linspace(0, 0.9, magma_r_cmap.N)))
projection = ccrs.LambertConformal(central_latitude=radar.latitude['data'][0],
                                   central_longitude=radar.longitude['data'][0])
# plot
plt.figure(figsize=(10,4))
ax1=plt.subplot(1,2,1, projection=projection)
display.plot_grid('reflectivity', vmin=5, vmax=45, cmap=ref_cmap, projection=projection,
                  transform=ccrs.PlateCarree(), ax=ax1, axislabels_flag=False)
ax2=plt.subplot(1,2,2, projection=projection)
display.plot_grid('convsf', vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3), projection=projection,
                  axislabels_flag=False, transform=ccrs.PlateCarree(), ticks=[1/3, 1, 5/3],
                  ticklabs=['', 'Stratiform', 'Convective'], ax=ax2)
plt.show()

######################################
# **Winter storm example with image muting**
# Here is a final example of the convective stratiform classification using an example from a winter storm. Before
# doing the classification, we image mute the reflectivity to remove regions with melting or mixed precipitation. We
# then rescale the reflectivity to snow rate (Rasumussen et al. 2003). We recommend using a rescaled reflectivity
# to do the classification, but if you do make sure to changed dB_averaging to False because this parameter is  used
# to convert reflectivity to a linear value before averaging (set dB_averaging to True for reflectivity fields in
# dBZ units).
# In this example, note how we change some of the other parameters since we are classifying snow rate instead of
# reflecitivity.

# Read in file
nexrad_file = 's3://noaa-nexrad-level2/2021/02/07/KOKX/KOKX20210207_161413_V06'
radar = pyart.io.read_nexrad_archive(nexrad_file)

# extract the lowest sweep
radar = radar.extract_sweeps([0])

# interpolate to grid
grid = pyart.map.grid_from_radars(
    (radar,), grid_shape=(1, 201, 201),
    grid_limits=((0, 10000), (-200000.0, 200000.0), (-200000.0, 200000.0)),
    fields=['reflectivity', 'cross_correlation_ratio'])

# image mute grid object
grid = pyart.util.image_mute_radar(grid, 'reflectivity', 'cross_correlation_ratio', 0.97, 20)

# convect non-muted reflectivity to snow rate
nonmuted_ref = grid.fields['nonmuted_reflectivity']['data'][0,:,:]
nonmuted_ref = np.ma.masked_invalid(nonmuted_ref)

nonmuted_ref_linear = 10 ** (nonmuted_ref / 10)  # mm6/m3
snow_rate = (nonmuted_ref_linear/57.3)**(1/1.67) #

# add to grid
snow_rate_dict = {
        'data': snow_rate[None,:,:],
        'standard_name': 'snow_rate',
        'long_name': 'Snow rate converted from linear reflectivity',
        'units': 'mm/hr',
        'valid_min': 0,
        'valid_max': 40500}
grid.add_field('snow_rate', snow_rate_dict, replace_existing=True)

# get dx dy
dx = grid.x['data'][1] - grid.x['data'][0]
dy = grid.y['data'][1] - grid.y['data'][0]

# convective stratiform classification
convsf_dict = pyart.retrieve.conv_strat_yuter(grid, dx, dy, refl_field='snow_rate', dB_averaging=False,
                                            always_core_thres=4,bkg_rad_km=40, use_cosine=True, max_diff=1.5,
                                            zero_diff_cos_val=5,weak_echo_thres=0, min_dBZ_used=0,
                                            max_conv_rad_km=1, estimate_flag=False)

# add to grid object
# mask zero values (no surface echo)
convsf_masked = np.ma.masked_equal(convsf_dict['convsf']['data'], 0)
# mask 3 values (weak echo)
convsf_masked = np.ma.masked_equal(convsf_masked, 3)
# add dimension to array to add to grid object
convsf_dict['convsf']['data'] = convsf_masked[None,:,:]
# add field
grid.add_field('convsf', convsf_dict['convsf'], replace_existing=True)

# create plot using GridMapDisplay
# plot variables
display = pyart.graph.GridMapDisplay(grid)
magma_r_cmap = plt.get_cmap('magma_r')
ref_cmap = mcolors.LinearSegmentedColormap.from_list('ref_cmap', magma_r_cmap(np.linspace(0, 0.9, magma_r_cmap.N)))
projection = ccrs.LambertConformal(central_latitude=radar.latitude['data'][0],
                                   central_longitude=radar.longitude['data'][0])
# plot
plt.figure(figsize=(10,4))
ax1=plt.subplot(1,2,1, projection=projection)
display.plot_grid('snow_rate', vmin=0, vmax=10, cmap=plt.get_cmap('viridis'), projection=projection,
                  transform=ccrs.PlateCarree(), ax=ax1, axislabels_flag=False)
ax2=plt.subplot(1,2,2, projection=projection)
display.plot_grid('convsf', vmin=0, vmax=2, cmap=plt.get_cmap('viridis', 3), projection=projection,
                  axislabels_flag=False, transform=ccrs.PlateCarree(), ticks=[1/3, 1, 5/3],
                  ticklabs=['', 'Stratiform', 'Convective'], ax=ax2)
plt.show()

######################################
# **Summary of recommendations and best practices**
# - Tune your parameters to your specific purpose
# - Use a rescaled field if possible (i.e. linear reflectivity, rain or snow rate)
# - Keep `estimate_flag=True` to see uncertainty in classification