"""
STEPS (Hong Kong)
========================================================
This example demonstrates how to make use of STEPS to perform 3-hour rainfall
nowcasting with radar data in Hong Kong
"""

###########################################################
# 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 time calculations
import pandas as pd
# Python package for numerical calculations
import numpy as np
# Python package for xarrays to read and handle netcdf data
import xarray as xr
# Python package for text formatting
import textwrap
# Python package for projection description
from pyresample import get_area_def
# Python package for projection
import cartopy.crs as ccrs
# Python package for land/sea features
import cartopy.feature as cfeature
# Python package for reading map shape file
import cartopy.io.shapereader as shpreader
# Python package for creating plots
from matplotlib import pyplot as plt
# Python package for output grid format
from matplotlib.gridspec import GridSpec
# Python package for colorbars
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.cm import ScalarMappable

# swirlspy data parser function
from swirlspy.rad.iris import read_iris_grid
# swirlspy test data source locat utils function
from swirlspy.qpe.utils import timestamps_ending, locate_file
# swirlspy regrid function
from swirlspy.core.resample import grid_resample
# swirlspy standardize data function
from swirlspy.utils import FrameType, standardize_attr, conversion
# swirlspy pysteps integrated package
from swirlspy.qpf import steps, dense_lucaskanade
# directory constants
from swirlspy.tests.samples import DATA_DIR
from swirlspy.tests.outputs import OUTPUT_DIR

warnings.filterwarnings("ignore")

###########################################################
# Define the working directory and nowcast parameters

radar_dir = os.path.abspath(
    os.path.join(DATA_DIR, 'iris/ppi')
)

# Set nowcast parameters
n_timesteps = int(3 * 60 / 6)  # 3 hours, each timestamp is 6 minutes
n_ens_members = 3

########################################################################
# Define the User Grid

area_id = "hk1980_250km"
description = ("A 250 m resolution rectangular grid "
               "centred at HKO and extending to 250 km "
               "in each direction in HK1980 easting/northing coordinates")
proj_id = 'hk1980'
projection = ('+proj=tmerc +lat_0=22.31213333333334 '
              '+lon_0=114.1785555555556 +k=1 +x_0=836694.05 '
              '+y_0=819069.8 +ellps=intl +towgs84=-162.619,-276.959,'
              '-161.764,0.067753,-2.24365,-1.15883,-1.09425 +units=m '
              '+no_defs')
x_size = 500
y_size = 500
area_extent = (587000, 569000, 1087000, 1069000)
area_def_tgt = get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

###########################################################
# Define the plotting function:

# Defining plot parameters
map_shape_file = os.path.abspath(os.path.join(
    DATA_DIR,
    'shape/hk'
))

# coastline and province
map_with_province = cfeature.ShapelyFeature(
    list(shpreader.Reader(map_shape_file).geometries()),
    ccrs.PlateCarree()
)


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

    # 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, province and state, color
    ax.add_feature(map_with_province,
                   facecolor=cfeature.COLORS['land'], edgecolor='none', zorder=0)
    # overlay coastline, province and state without color
    ax.add_feature(map_with_province, facecolor='none',
                   edgecolor='gray', linewidth=0.5)

    ax.set_title('')

#############################################################
# Loading radar data
# ---------------------------------------------------


# Log the start time
start_time = pd.Timestamp.now()

# Define the basetime
basetime = pd.Timestamp('201902190800')

# Generate the timestamps and locate the files
located_files = []
radar_ts = timestamps_ending(
    basetime,
    duration=pd.Timedelta(minutes=12),
    exclude_end=False
)
for timestamp in radar_ts:
    located_files.append(locate_file(radar_dir, timestamp))

# Read in the data
reflectivity_list = []  # stores reflec from read_iris_grid()
for filename in located_files:
    reflec = read_iris_grid(filename)
    reflectivity_list.append(reflec)

# Reproject the radar data to the user-defined grid
area_def_src = reflectivity_list[0].attrs['area_def']
reproj_reflectivity_list = []
for reflec in reflectivity_list:
    reproj_reflec = grid_resample(
        reflec, area_def_src, area_def_tgt,
        coord_label=['x', 'y']
    )
    reproj_reflectivity_list.append(reproj_reflec)

# Fill in all fields of the xarray of reflectivity data
frames = xr.concat(reproj_reflectivity_list,
                   dim='time').sortby(['y'], ascending=False)
standardize_attr(frames, frame_type=FrameType.dBZ)


# Convert from reflectivity to rainfall rain
frames = conversion.to_rainfall_rate(frames, True, a=58.53, b=1.56)

# Set the fill value
frames.attrs['zero_value'] = -15.0

# apply threshold to -10dBR i.e. 0.1mm/h
threshold = -10.0
frames.values[frames.values < threshold] = frames.attrs['zero_value']

# Set missing values with the fill value
frames.values[~np.isfinite(frames.values)] = frames.attrs['zero_value']

# Log the time for record
initialising_time = pd.Timestamp.now()

###########################################################
# Running Lucas Kanade Optical flow and S-PROG
# -------------------------------------------
#

# Estimate the motion field with Lucas Kanade
motion = dense_lucaskanade(frames)

# Log the time for record
motion_time = pd.Timestamp.now()

# Nowcast using STEP
forcast_frames = steps(
    frames,
    motion,
    n_timesteps,
    n_ens_members=n_ens_members,
    n_cascade_levels=8,
    R_thr=threshold,
    kmperpixel=2.0,
    decomp_method="fft",
    bandpass_filter_method="gaussian",
    noise_method="nonparametric",
    probmatching_method="mean",
    vel_pert_method="bps",
    mask_method="incremental",
    seed=24
)

steps_time = pd.Timestamp.now()

#############################################
# Generating radar reflectivity maps
# -----------------------------------
#

# Defining the colour scale
levels = [
    -32768,
    10, 15, 20, 24, 28, 32,
    34, 38, 41, 44, 47, 50,
    53, 56, 58, 60, 62
]
cmap = ListedColormap([
    '#FFFFFF00', '#08C5F5', '#0091F3', '#3898FF', '#008243', '#00A433',
    '#00D100', '#01F508', '#77FF00', '#E0D100', '#FFDC01', '#EEB200',
    '#F08100', '#F00101', '#E20200', '#B40466', '#ED02F0'
])

norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

mappable = ScalarMappable(cmap=cmap, norm=norm)
mappable.set_array([])

# Defining the crs
crs = area_def_tgt.to_cartopy_crs()

# Defining area
x = frames.coords['x'].values
y = frames.coords['y'].values
extents = [
    x[0], y[0],
    x[-1], y[-1]
]

# Generate a time steps for every hour
time_steps = [
    (basetime + pd.Timedelta(minutes=6*i)) 
    for i in range(n_timesteps + 1) if i % 10 == 0
]

ref_frames = conversion.to_reflectivity(forcast_frames, True)
ref_frames.data[ref_frames.data < 0.1] = np.nan

qx = motion.coords['x'].values[::25]
qy = motion.coords['y'].values[::25]
qu = motion.values[0, ::25, ::25]
qv = motion.values[1, ::25, ::25]

n_rows = len(time_steps)
fig: plt.Figure = plt.figure(
    figsize=(n_ens_members * 3.5 + 1, n_rows * 3.5), frameon=False)
gs = GridSpec(n_rows, n_ens_members, figure=fig,
              wspace=0.03, hspace=0, top=0.95, bottom=0.05, left=0.17, right=0.845)

for row in range(n_rows):
    for col in range(n_ens_members):
        ax: plt.Axes = fig.add_subplot(gs[row, col], projection=crs)

        ensemble = ref_frames.coords['ensembles'].values[col]
        t = time_steps[row]

        # plot base map
        plot_base(ax, extents, crs)

        # plot reflectivity
        member = ref_frames.sel(ensembles=ensemble)
        frame = member.sel(time=t)
        im = ax.imshow(frame.values, cmap=cmap, norm=norm, interpolation='nearest',
                       extent=extents)

        # plot motion vector
        ax.quiver(
            qx, qy, qu, qv, pivot='mid'
        )

        ax.text(
            extents[0],
            extents[1],
            textwrap.dedent(
                """
                    Reflectivity
                    Based @ {baseTime}
                    """
            ).format(
                baseTime=basetime.strftime('%H:%MH')
            ).strip(),
            fontsize=10,
            va='bottom',
            ha='left',
            linespacing=1
        )
        ax.text(
            extents[2] - (extents[2] - extents[0]) * 0.03,
            extents[1],
            textwrap.dedent(
                """
                    {validDate}
                    Valid @ {validTime}
                    """
            ).format(
                validDate=basetime.strftime('%Y-%m-%d'),
                validTime=t.strftime('%H:%MH')
            ).strip(),
            fontsize=10,
            va='bottom',
            ha='right',
            linespacing=1
        )

cbar_ax = fig.add_axes([0.875, 0.075, 0.03, 0.85])
cbar = fig.colorbar(
    mappable, cax=cbar_ax, ticks=levels[1:], extend='max', format='%.3g')
cbar.ax.set_ylabel(ref_frames.attrs['values_name'], rotation=90)


fig.savefig(
    os.path.join(
        OUTPUT_DIR,
        "steps-reflectivity.png"
    ),
    bbox_inches='tight'
)

radar_image_time = pd.Timestamp.now()

##########################################################################
# Accumulating hourly rainfall for 3-hour forecast
# ------------------------------------------------
#
# Hourly accumulated rainfall is calculated every 30 minutes, the first
# endtime is the basetime i.e. T+30min.
#

# Convert from rainfall rate to rainfall
rf_frames = conversion.to_rainfall_depth(ref_frames, a=58.53, b=1.56)

# Compute hourly accumulated rainfall every 30 minutes.
acc_rf_frames = []
for ens in rf_frames.coords['ensembles']:
    af = conversion.acc_rainfall_depth(
        rf_frames.sel(ensembles=ens).drop('ensembles'), basetime +
        pd.Timedelta(minutes=60), basetime + pd.Timedelta(hours=3)
    )
    af_ensembles = af.assign_coords(ensembles=ens)
    acc_rf_frames.append(af_ensembles.expand_dims('ensembles'))
acc_rf_frames = xr.concat(acc_rf_frames, dim='ensembles')

# Replace zero value with NaN
acc_rf_frames.data[acc_rf_frames.data <=
                   acc_rf_frames.attrs['zero_value']] = np.nan

acc_time = pd.Timestamp.now()

#############################################
# Generating radar reflectivity maps
# -----------------------------------
#

# Defining colour scale
levels = [
    0, 0.5, 2, 5, 10, 20,
    30, 40, 50, 70, 100, 150,
    200, 300, 400, 500, 600, 700
]
cmap = ListedColormap([
    '#FFFFFF00', '#9bf7f7', '#00ffff', '#00d5cc', '#00bd3d', '#2fd646',
    '#9de843', '#ffdd41', '#ffac33', '#ff621e', '#d23211', '#9d0063',
    '#e300ae', '#ff00ce', '#ff57da', '#ff8de6', '#ffe4fd'
])

norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

mappable = ScalarMappable(cmap=cmap, norm=norm)
mappable.set_array([])

time_steps = acc_rf_frames.coords['time'].values
n_rows = len(time_steps)

fig: plt.Figure = plt.figure(
    figsize=(n_ens_members * 4 + 1, n_rows * 4), frameon=False)
gs = GridSpec(n_rows, n_ens_members, figure=fig,
              wspace=0.03, hspace=-0.25, top=0.95, bottom=0.05, left=0.17, right=0.845)

for row in range(n_rows):
    for col in range(n_ens_members):
        ax = fig.add_subplot(gs[row, col], projection=crs)

        ensemble = acc_rf_frames.coords['ensembles'].values[col]
        t = time_steps[row]

        # plot base map
        plot_base(ax, extents, crs)

        # plot accumulated rainfall depth
        t = pd.Timestamp(t)
        member = acc_rf_frames.sel(ensembles=ensemble)
        frame = member.sel(time=t)
        im = ax.imshow(frame.values, cmap=cmap, norm=norm, interpolation='nearest',
                       extent=extents)

        ax.text(
            extents[0],
            extents[1],
            textwrap.dedent(
                """
                    Hourly Rainfall
                    Based @ {baseTime}
                    """
            ).format(
                baseTime=basetime.strftime('%H:%MH')
            ).strip(),
            fontsize=10,
            va='bottom',
            ha='left',
            linespacing=1
        )
        ax.text(
            extents[2] - (extents[2] - extents[0]) * 0.03,
            extents[1],
            textwrap.dedent(
                """
                    {validDate}
                    Valid @ {validTime}
                    """
            ).format(
                validDate=basetime.strftime('%Y-%m-%d'),
                validTime=t.strftime('%H:%MH')
            ).strip(),
            fontsize=10,
            va='bottom',
            ha='right',
            linespacing=1
        )

cbar_ax = fig.add_axes([0.875, 0.095, 0.03, 0.8])
cbar = fig.colorbar(
    mappable, cax=cbar_ax, ticks=levels[1:], extend='max', format='%.3g')
cbar.ax.set_ylabel(acc_rf_frames.attrs['values_name'], rotation=90)

fig.savefig(
    os.path.join(
        OUTPUT_DIR,
        "steps-rainfall.png"
    ),
    bbox_inches='tight'
)

ptime = pd.Timestamp.now()

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

print(f"Start time: {start_time}")
print(f"Initialising time: {initialising_time}")
print(f"Motion field time: {motion_time}")
print(f"STEPS time: {steps_time}")
print(f"Plotting radar image time: {radar_image_time}")
print(f"Accumulating rainfall time: {acc_time}")
print(f"Plotting rainfall maps: {ptime}")

print(f"Time to initialise: {initialising_time - start_time}")
print(f"Time to run motion field: {motion_time - initialising_time}")
print(f"Time to perform STEPS: {steps_time - motion_time}")
print(f"Time to plot radar image: {radar_image_time - steps_time}")
print(f"Time to accumulate rainfall: {acc_time - radar_image_time}")
print(f"Time to plot rainfall maps: {ptime - acc_time}")

print(f"Total: {ptime - start_time}")