"""
Simple Lightning Nowcast (Hong Kong)
========================================================
This example demonstrates how to perform
a simple lightning nowcast by the extrapolation
of lightning strikes, using data from 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
import pyproj
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

# Python com-swirls package to calculate motion field (rover) and semi-lagrangian advection
from swirlspy.qpf import rover, sla
# 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 lightning function
from swirlspy.obs import Lightning
from swirlspy.ltg.map import gen_ltg_binary
# swirlspy standardize data function
from swirlspy.utils import standardize_attr, FrameType
# directory constants
from swirlspy.tests.samples import DATA_DIR
from swirlspy.tests.outputs import OUTPUT_DIR

warnings.filterwarnings("ignore")
plt.switch_backend('agg')

###########################################################################
# Initialising
# -------------------------------------------------------------------------
#
# Define the target grid:
#

# Defining target grid
area_id = "hk1980_250km"
description = ("A 1km 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 = get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

############################################################################
# Define basemap:

# Load the shape of Hong Kong
map_shape_file = os.path.abspath(os.path.join(
    DATA_DIR,
    'shape/hk_hk1980'
))

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


# define the plot function
def plot_base(ax: plt.Axes, extents: list, crs: ccrs.Projection):
    # fake the ocean color
    ax.imshow(
        np.tile(np.array([[[178, 208, 254]]], dtype=np.uint8), [2, 2, 1]),
        origin='upper', transform=crs,
        extent=extents, 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('')


############################################################################
# Defining a basetime, nowcast interval and nowcast duration:
#

basetime = pd.Timestamp('201904201200').floor('min')
basetime_str = basetime.strftime('%Y%m%d%H%M')
nowcast_interval = pd.Timedelta(30, 'm')
nowcast_duration = pd.Timedelta(2, 'h')

############################################################################
# Loading Lightning and Radar Data
# --------------------------------------------------------------------------
#
# Extract Lightning Location Information System (LLIS) files to read.
#

# The interval between LLIS files
llis_file_interval = pd.Timedelta(1, 'm')

# Finding the string representation
# of timestamps marking the relevant LLIS files
timestamps = timestamps_ending(
    basetime,
    duration=nowcast_interval,
    interval=llis_file_interval,
    format='%Y%m%d%H%M'
)

# Locating files
llis_dir = os.path.join(DATA_DIR, 'llis')

located_files = []
for timestamp in timestamps:
    located_file = locate_file(llis_dir, timestamp)
    located_files.append(located_file)
    if located_file is not None:
        print(f"LLIS file for {timestamp} has been located")

###################################################################
# Read data from files into Lightning objects.
#

# Initialising Lightning objects
# Coordinates are geodetic
ltg_object_geodetic = Lightning(
    located_files,
    'geodetic',
    basetime - nowcast_interval,
    basetime
)

#####################################################################
# Locating IRIS REF radar files needed to generate motion field.
# The IRIS REF files are marked by start time.
#

# The interval between the radar files
rad_interval = pd.Timedelta(6, 'm')

# Finding the string representation
# of the timestamps
# marking the required files
# In this case, the most recent radar files
# separated by `nowcast_interval` is required

rad_timestamps = timestamps_ending(
    basetime,
    duration=nowcast_interval + rad_interval,
    interval=rad_interval
)

rad_timestamps = [rad_timestamps[0], rad_timestamps[-1]]

rad_dir = os.path.join(DATA_DIR, 'iris')
located_rad_files = []
for timestamp in rad_timestamps:
    located_file = locate_file(rad_dir, timestamp)
    located_rad_files.append(located_file)
    if located_file is not None:
        print(f"IRIS REF file for {timestamp} has been located")

##########################################################################
# Read from iris radar files using read_iris_grid().
# The data is in the Azimuthal Equidistant Projection.
#

reflec_unproj_list = []
for file in located_rad_files:
    reflec = read_iris_grid(file)
    reflec_unproj_list.append(reflec)


#############################################################################
# Data Reprojection
# ---------------------------------------------------------------------------
#
# Since the coordinates of the Lightning objects are geodetic, a conversion
# from geodetic coordinates to that of the user-defined projection system is
# required.
# This can be achieved by `swirlspy.obs.Lightning.reproject()`
#

# pyproj representation of the map projection of the input data
inProj = pyproj.Proj(init='epsg:4326')
# pyproj representation of the intended map projection
outProj = pyproj.Proj(area_def.proj_str)

# reproject
ltg_object = ltg_object_geodetic.reproject(inProj, outProj, 'HK1980')

#########################################################################
# Radar data also needs to be reprojected from Azimuthal Equidistant
# to HK1980. This is achieved by the `swirlspy.core.resample` function.
# Following reprojection, the individual xarray.DataArrays can be concatenated
# to a single xarray.DataArrays.
#

reflec_list = []
for z in reflec_unproj_list:
    reflec = grid_resample(
        z,
        z.attrs['area_def'],
        area_def,
        coord_label=['easting', 'northing']
    )
    reflec_list.append(reflec)

reflectivity = xr.concat(reflec_list, dim='time')
standardize_attr(reflectivity, frame_type=FrameType.dBZ)

print(reflectivity)

####################################################################
# Grid lightning observations
# -----------------------------------------------------------------
#
# Generate a binary grid of lighting observations, where 0
# indicates no lightning and 1 indicates that there is lightning.
#

ltg_bin = gen_ltg_binary(ltg_object,
                         area_def,
                         coord_label=['easting', 'northing'])
print(ltg_bin)

# Duplicating ltg_density along the time dimension
# so that it can be advected.
ltg_bin1 = ltg_bin.copy()
ltg_bin1.coords['time'] = [basetime - nowcast_interval]
ltg_bin_concat = xr.concat(
    [ltg_bin1, ltg_bin],
    dim='time'
)
# Identifying zero value
ltg_bin_concat.attrs['zero_value'] = np.nan


####################################################################
# Generating motion fields and extrapolating
# --------------------------------------------------------------
#
# #. Generate motion field using ROVER.
# #. Extrapolate gridded lightning strikes using the motion field
#    by Semi-Lagrangian Advection.
#

motion = rover(reflectivity)  # rover

steps = int(nowcast_duration.total_seconds()
            // nowcast_interval.total_seconds())
ltg_frames = sla(ltg_bin_concat, motion, nowcast_steps=steps-1)  # sla
print(ltg_frames)


##################################################################
# Visualization
# ---------------------------------------------------------------
#
# Now, let us visualize the lightning nowcast up to a duration
# of 2 hours.
#

# Colormap and colorbar parameters
n_col = 2
cmap = ListedColormap(['#FFFFFF', '#000000'])
levels = [0., 0.5, 1.]
ticks = np.array([0, 1])
norm = BoundaryNorm(levels, cmap.N, clip=True)

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

# Defining the crs
crs = area_def.to_cartopy_crs()

# Defining area
x = ltg_frames.coords['easting'].values
y = ltg_frames.coords['northing'].values
extents = [
    np.min(x), np.max(x),
    np.min(y), np.max(y)
]

qx = motion.coords['easting'].values[::20]
qy = motion.coords['northing'].values[::20]
qu = motion.values[0, ::20, ::20]
qv = motion.values[1, ::20, ::20]

fig = plt.figure(figsize=(8, 8), frameon=False)
gs = GridSpec(
    2, 2, figure=fig, wspace=0.03, hspace=-0.25, top=0.95,
    bottom=0.05, left=0.17, right=0.845
)

for i, t in enumerate(ltg_frames.time.values):
    time = pd.Timestamp(t)

    row = i // 2
    col = i % 2
    ax = fig.add_subplot(gs[row, col], projection=crs)

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

    # plot lightning
    frame = ltg_frames.sel(time=time)
    frame.where(frame > levels[1]).plot(
        cmap=cmap,
        norm=norm,
        add_colorbar=False
    )

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

    ax.set_title('')

    ax.text(
        extents[0],
        extents[3],
        textwrap.dedent(
            """
            Lightning Density
            Based @ {baseTime}
            """
        ).format(
            baseTime=basetime.strftime('%H:%MH')
        ).strip(),
        fontsize=10,
        va='bottom',
        ha='left',
        linespacing=1
    )

    ax.text(
        extents[1],
        extents[3],
        textwrap.dedent(
            """
            {validDate}
            Valid @ {validTime}
            """
        ).format(
            validDate=basetime.strftime('%Y-%m-%d'),
            validTime=time.strftime('%H:%MH')
        ).strip(),
        fontsize=10,
        va='bottom',
        ha='right',
        linespacing=1
    )

cbar_ax = fig.add_axes([0.875, 0.125, 0.03, 0.75])
cbar = fig.colorbar(mappable, cax=cbar_ax)
tick_locs = (np.arange(n_col) + 0.5)*(n_col-1)/n_col
cbar.set_ticks(tick_locs)
cbar.set_ticklabels(np.arange(n_col))
cbar.ax.set_ylabel(
    f"{ltg_bin.attrs['long_name']}[{ltg_bin.attrs['units']}]"
)

fig.savefig(
    os.path.join(OUTPUT_DIR, "ltg-bin-hk.png"),
    bbox_inches='tight'
)