""" 
Radar Advection with NWP Blend (Malaysia) 
===============================================================
This example shows the technique of blending Numerical Weather 
Forecast (NWP) with COM-SWIRLS output to generate operational 
deterministic QPF up to 3 and 6 hours ahead. NWP data was generated
using the WRF Model while radar data was taken from the 
Malaysian radar observation network

"""

###########################################################
# Definitions
# --------------------------------------------------------
#

########################################################################
# Import all required modules and methods:

# Python package to allow system command line functions
import os
# Python package to manage warning message
import warnings
# Python package for timestamp
import pandas as pd
# Python package for xarrays to read and handle netcdf data
import xarray as xr
# Python package for numerical calculations
import numpy as np
# Python package for reading map shape file
import cartopy.io.shapereader as shpreader
# Python package for land/sea features
import cartopy.feature as cfeature
# Python package for projection
import cartopy.crs as ccrs
# Python package for creating plots
from matplotlib import pyplot as plt
# Python package for output import grid
from matplotlib.gridspec import GridSpec
# Python package for colorbars
from matplotlib.colors import BoundaryNorm, ListedColormap
# Python package for scalar data to RGBA mapping
from matplotlib.cm import ScalarMappable
# Python com-swirls package to standardize attributes
from swirlspy.utils import standardize_attr, FrameType
# Python com-swirls package to calculate motion field (rover) and semi-lagrangian advection
from swirlspy.qpf import rover, sla
# Python com-swirls package to blend nwp and nowcast (RaINS)
from swirlspy.blending import rains, nwp_bias_correction
# directory constants
from swirlspy.tests.samples import DATA_DIR
from swirlspy.tests.outputs import OUTPUT_DIR

warnings.filterwarnings("ignore")
start_time = pd.Timestamp.now()

#############################################################
# Initializing
# ---------------------------------------------------
#
# This section demonstrates the extraction of radar & nwp data
# from netcdf into python
#
# Step 1: Define your input data directory and output directory
#

# Supply the directory of radar and nwp data
data_dir = os.path.abspath(
    os.path.join(DATA_DIR, 'netcdf_ms')
)

##############################################################################
# Step 2: Define a basetime
#

# Supply basetime
basetime = pd.Timestamp('201908090900')

##############################################################################
# Step 3: Read data files from the radar data using xarray()
#

# Radar data listed from the basetime[0] --> 3 hours before the basetime[17] (descending time)
interval = 10             # Interval of radar data
radar_datas = []
for i in range(0, 18):
    t = basetime - pd.Timedelta(minutes=i * interval)
    # Radar data nomenclature
    filename = os.path.join(
        data_dir,
        t.strftime("radar_d03_%Y-%m-%d_%H_%M_00.rapids.nc")
    )
    reflec = xr.open_dataset(filename)
    radar_datas.append(reflec)

# Concatenate list by time
reflec_concat = xr.concat(radar_datas, dim='time')

# Extracting the radar data: The radar dBZ variable is named 'Zradar', therefore, we extract 'Zradar'
radar = reflec_concat['Zradar']

# Reversing such that time goes from earliest to latest; 3 hours before basetime[0] --> basetime[17]
radar = radar.sortby('time', ascending=True)

##############################################################################
# Step 4: Reading nwp netcdf data into xarray
#

# NWP data listed from the basetime[0] --> 3 and 6 hours AFTER basetime[17] (nowcast)
nwp_3hr_datas = []
nwp_6hr_datas = []
for i in range(0, 36):
    t = basetime + pd.Timedelta(minutes=i * interval)
    # Radar nwp nomenclature
    filename = os.path.join(
        data_dir,
        t.strftime("wrfout_d03_%Y-%m-%d_%H_%M_00.rapids.nc")
    )
    reflec = xr.open_dataset(filename)
    nwp_6hr_datas.append(reflec)
    if i < 18:
        nwp_3hr_datas.append(reflec)

# Concatenating the nwp reflectivity list of data by time
reflec_3hr_concat = xr.concat(nwp_3hr_datas, dim='time')
reflec_6hr_concat = xr.concat(nwp_6hr_datas, dim='time')

# Extracting the nwp data: The nwp dBZ variable is called 'zwrf', extracting 'zwrf'
nwp_3hr = reflec_3hr_concat['zwrf']
nwp_6hr = reflec_6hr_concat['zwrf']

initialising_time = pd.Timestamp.now()

##############################################################################
# Step 5: Bias correction of nwp data
# The objective of bias correction is to match the nwp percentile to the radar percentile
# This is also known as frequency matching.
#

nwp_3hr_corrected = nwp_bias_correction(radar, nwp_3hr)
nwp_6hr_corrected = nwp_bias_correction(radar, nwp_6hr)

bias_correction_time = pd.Timestamp.now()

##############################################################################
# Nowcast (SWIRLS-Radar-Advection)
# ---------------------------------------------------
# The swirls radar advection was performed using the observed radar data
# Firstly, some attributes necessary for com-swirls input variable is added
# Secondly, rover function is invoked to compute the motion field
# Thirdly, semi-lagrangian advection is performed to advect the radar data using the rover motion field

# Adding in some attributes that is step_size <10 mins in pandas.Timedelta>, zero_value <9999.> frame_type <FrameType.dBZ>
radar.attrs['step_size'] = pd.Timedelta(minutes=10)
standardize_attr(radar, frame_type=FrameType.dBZ, zero_value=np.nan)

# Rover motion field computation
motion = rover(radar)

rover_time = pd.Timestamp.now()

# Semi-Lagrangian Advection
swirls = sla(radar, motion, 35)  # Radar time goes from earliest to latest

sla_time = pd.Timestamp.now()

##############################################################################
# Blending (RaINS)
# ---------------------------------------------------
# Blending Numerical Weather Forecast (NWP) with COM-SWIRLS output to generate operational
# deterministic QPF up to 3 hours ahead.

blended_3hr = rains(nwp_3hr_corrected, swirls[0:18])
blended_6hr = rains(nwp_6hr_corrected, swirls)

rains_time = pd.Timestamp.now()

##############################################################################
# Plotting result
# ---------------------------------------------------
# Step 1: Defining the dBZ levels, colorbar parameters and projection

# levels of colorbar (dBZ)
levels = [-32768, 10, 15, 20, 24, 28, 32, 34, 38, 41, 44,
          47, 50, 53, 56, 58, 60, 62]
# hko colormap for dBZ at each levels
cmap = ListedColormap([
    '#FFFFFF', '#08C5F5', '#0091F3', '#3898FF', '#008243', '#00A433',
    '#00D100', '#01F508', '#77FF00', '#E0D100', '#FFDC01', '#EEB200',
    '#F08100', '#F00101', '#E20200', '#B40466', '#ED02F0'
])
# boundary
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)
# scalar data to RGBA mapping
scalar_map = ScalarMappable(cmap=cmap, norm=norm)
scalar_map.set_array([])
# Defining plot parameters
map_shape_file = os.path.abspath(os.path.join(
    DATA_DIR,
    'shape/se_asia_s'
))
# coastline and province
se_asia = cfeature.ShapelyFeature(
    list(shpreader.Reader(map_shape_file).geometries()),
    ccrs.PlateCarree()
)
# output area
extents = [99, 120, 0.5, 7.25]

# base_map plotting function


def plot_base(ax: plt.Axes):
    ax.set_extent(extents, crs=ccrs.PlateCarree())

    # fake the ocean color
    ax.imshow(np.tile(np.array([[[178, 208, 254]]],
                               dtype=np.uint8), [2, 2, 1]),
              origin='upper',
              transform=ccrs.PlateCarree(),
              extent=[-180, 180, -180, 180],
              zorder=-1)
    # coastline, state, color
    ax.add_feature(se_asia,
                   facecolor=cfeature.COLORS['land'], edgecolor='none', zorder=0)
    # overlay coastline, state without color
    ax.add_feature(se_asia, facecolor='none',
                   edgecolor='gray', linewidth=0.5)

    ax.set_title('')


##############################################################################
# Step 2: Filtering
# values <= 5dbZ are not plotted
nwp_3hr_corrected = nwp_3hr_corrected.where(nwp_3hr_corrected > 5, np.nan)
blended_3hr = blended_3hr.where(blended_3hr > 5, np.nan)
nwp_6hr_corrected = nwp_6hr_corrected.where(nwp_6hr_corrected > 5, np.nan)
blended_6hr = blended_6hr.where(blended_6hr > 5, np.nan)

##############################################################################
# Step 3: Plotting the swirls-radar-advection, nwp-bias-corrected, blended 3 hours ahead
#
fig: plt.Figure = plt.figure(
    figsize=(3 * 5 + 1, 3 * 2),
    frameon=False
)
gs = GridSpec(
    3, 3, figure=fig,
    wspace=0.03, hspace=0, top=0.95, bottom=0.05, left=0.17, right=0.845
)

for row in range(3):
    time_index = (row + 1) * 6 - 1
    timelabel = basetime + pd.Timedelta(interval * (time_index + 1), 'm')

    for col in range(3):
        ax: plt.Axes = fig.add_subplot(
            gs[row, col],
            projection=ccrs.PlateCarree()
        )
        if col % 3 == 0:
            z = swirls[time_index].values
            lats = swirls[time_index].latitude
            lons = swirls[time_index].longitude
            title = 'SWIRLS Reflectivity'
        elif col % 3 == 1:
            z = nwp_3hr_corrected[time_index].values
            lats = nwp_3hr_corrected[time_index].lat
            lons = nwp_3hr_corrected[time_index].lon
            title = 'NWP Reflectivity'
        elif col % 3 == 2:
            z = blended_3hr[time_index].values
            lats = blended_3hr[time_index].latitude
            lons = blended_3hr[time_index].longitude
            title = 'Blended Reflectivity'

        # plot base map
        plot_base(ax)

        # plot reflectivity
        ax.contourf(
            lons, lats, z, 60,
            transform=ccrs.PlateCarree(),
            cmap=cmap, norm=norm, levels=levels
        )

        ax.set_title(
            f"{title}\n" +
            f"Initial @ {basetime.strftime('%H:%MZ')}",
            loc='left', fontsize=9
        )
        ax.set_title('')
        ax.set_title(
            f"Initial {basetime.strftime('%Y-%m-%d')} \n" +
            f"Forecast Valid @ {timelabel.strftime('%H:%MZ')} ",
            loc='right', fontsize=9
        )

cbar_ax = fig.add_axes([0.85, 0.105, 0.01, 0.845])
cbar = fig.colorbar(
    scalar_map, cax=cbar_ax, ticks=levels[1:], extend='max', format='%.3g'
)
cbar.ax.set_ylabel('Reflectivity (dBZ)', rotation=90)

fig.savefig(
    os.path.join(
        OUTPUT_DIR,
        "swirls_nwp_blend_ms_3hr.png"
    ),
    dpi=450,
    bbox_inches="tight",
    pad_inches=0.1
)

##############################################################################
# Step 4: Plotting the swirls-radar-advection, nwp-bias-corrected, blended 6 hours ahead
#
fig: plt.Figure = plt.figure(
    figsize=(3 * 5 + 1, 6 * 2),
    frameon=False
)
gs = GridSpec(
    6, 3, figure=fig,
    wspace=0.03, hspace=0, top=0.95, bottom=0.05, left=0.17, right=0.845
)

for row in range(6):
    time_index = (row + 1) * 6 - 1
    timelabel = basetime + pd.Timedelta(interval * (time_index + 1), 'm')

    for col in range(3):
        ax: plt.Axes = fig.add_subplot(
            gs[row, col],
            projection=ccrs.PlateCarree()
        )
        if col % 3 == 0:
            z = swirls[time_index].values
            lats = swirls[time_index].latitude
            lons = swirls[time_index].longitude
            title = 'SWIRLS Reflectivity'
        elif col % 3 == 1:
            z = nwp_6hr_corrected[time_index].values
            lats = nwp_6hr_corrected[time_index].lat
            lons = nwp_6hr_corrected[time_index].lon
            title = 'NWP Reflectivity'
        elif col % 3 == 2:
            z = blended_6hr[time_index].values
            lats = blended_6hr[time_index].latitude
            lons = blended_6hr[time_index].longitude
            title = 'Blended Reflectivity'

        # plot base map
        plot_base(ax)

        # plot reflectivity
        ax.contourf(
            lons, lats, z, 60,
            transform=ccrs.PlateCarree(),
            cmap=cmap, norm=norm, levels=levels
        )

        ax.set_title(
            f"{title}\n" +
            f"Initial @ {basetime.strftime('%H:%MZ')}",
            loc='left', fontsize=9
        )
        ax.set_title('')
        ax.set_title(
            f"Initial {basetime.strftime('%Y-%m-%d')} \n" +
            f"Forecast Valid @ {timelabel.strftime('%H:%MZ')} ",
            loc='right', fontsize=9
        )

cbar_ax = fig.add_axes([0.85, 0.105, 0.01, 0.845])
cbar = fig.colorbar(
    scalar_map, cax=cbar_ax, ticks=levels[1:], extend='max', format='%.3g'
)
cbar.ax.set_ylabel('Reflectivity (dBZ)', rotation=90)

fig.savefig(
    os.path.join(
        OUTPUT_DIR,
        "swirls_nwp_blend_ms_6hr.png"
    ),
    dpi=450,
    bbox_inches="tight",
    pad_inches=0.1
)

radar_image_time = pd.Timestamp.now()

######################################################################
# Checking run time of each component
# --------------------------------------------------------------------
#

print(f"Start time: {start_time}")
print(f"Initialising time: {initialising_time}")
print(f"NWP bias correction time: {bias_correction_time}")
print(f"SLA time: {sla_time}")
print(f"RaINS time: {rains_time}")
print(f"Plotting radar image time: {radar_image_time}")

print(f"Time to initialise: {initialising_time - start_time}")
print(
    f"Time to run NWP bias correction: {bias_correction_time - initialising_time}")
print(f"Time to run rover: {rover_time - bias_correction_time}")
print(f"Time to perform SLA: {sla_time - rover_time}")
print(f"Time to perform RaINS: {rains_time - sla_time}")
print(f"Time to plot radar image: {radar_image_time - rains_time}")