"""
Probabilistic Lightning Density Nowcast (Hong Kong)
========================================================
This example demonstrates how to perform
probabilistic lightning nowcast by advecting
lightning density, using data from Hong Kong.
"""

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

import os
from copy import copy
import textwrap
import pyproj
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.gridspec import GridSpec
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from cartopy.io import shapereader
from matplotlib.cm import ScalarMappable

from pyresample import utils

from swirlspy.qpf import rover
from swirlspy.qpf import sla
from swirlspy.qpe.utils import timestamps_ending, locate_file
from swirlspy.rad.iris import read_iris_grid
from swirlspy.core.resample import grid_resample
from swirlspy.obs import Lightning
from swirlspy.ltg.map import gen_ltg_density
from swirlspy.utils import standardize_attr, FrameType

plt.switch_backend('agg')
THIS_DIR = os.getcwd()

###########################################################################
# 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 = utils.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(
    THIS_DIR,
    '../tests/samples/shape/hk_hk1980'
))

# coastline and province
map_with_province = cfeature.ShapelyFeature(
    list(shapereader.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(THIS_DIR, '../tests/samples/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(THIS_DIR, '../tests/samples/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)

####################################################################
# Generating lightning density field
# ------------------------------------------------------------------
#
# Generate a lightning density field.
#
# #. Count the number of lightning strikes in
#    a circular search area of 8km surrounding a pixel over
#    the past 30 minutes. Do this for all the pixels.
# #. This value is then spatially and temporally averaged to
#    give a lightning density, an indicator of lightning activity given as the
#    number of strikes/km\ :sup:`2`/hr.
#

# Obtaining lightning density from lightning objects
ltg_density = gen_ltg_density(
    ltg_object,
    area_def,
    ltg_object.start_time,
    ltg_object.end_time,
    tolerance=8000,
    coord_label=['easting', 'northing']
    )

print(ltg_density)

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

################################################################
# Generating motion fields and extrapolating
# --------------------------------------------------------------
#
# #. Generate motion field using ROVER.
# #. Extrapolate lightning density using the motion field
#    by Semi-Lagrangian Advection.
# #. Repeat using different parameters to form an ensemble.

# Different ROVER parameters
params = pd.DataFrame()
params['smallest_scale'] = [1, 1, 1, 2]
params['largest_scale'] = [7, 7, 7, 7]
params['rho'] = [9, 9, 9, 9]
params['alpha'] = [2000, 2000, 10000, 2000]
params['sigma'] = [1.5, 2.5, 1.5, 2.5]
params['tracking_interval'] = [30, 30, 30, 30]

xtrp_frames = []
for i in range(len(params)):
    # Generating motion field using ROVER
    motion = rover(
        reflectivity,
        start_level=params['smallest_scale'][i],
        max_level=params['largest_scale'][i],
        rho=params['rho'][i],
        alpha=params['alpha'][i],
        sigma=params['sigma'][i]
    )
    # SLA
    steps = int(nowcast_duration.total_seconds()
                // nowcast_interval.total_seconds())
    ltg_density_xtrp = sla(ltg_density_concat, motion, nowcast_steps=steps)
    xtrp_frames.append(ltg_density_xtrp)


# Create coordinates for `ensemble` dimension
ens_coords = ["Member-{:02}".format(i) for i in range(len(xtrp_frames))]

# Concat ltg_frames along the ensemble dimension
ltg_frames = xr.concat(xtrp_frames, pd.Index(ens_coords, name='ensemble'))

print(ltg_frames)

###################################################################
# Now, let us visualise the extrapolated lightning density fields
# for the next 30 minutes.

levels = [
    0, 0.125, 0.25, 0.5, 1,
    2, 4, 8, 16, 32, 999
]
cmap = ListedColormap([
    '#FFFFFF', '#5050FF', '#00FFFF', '#00C030', '#A0FF80',
    '#FFFF40', '#FF9830', '#FF0000', '#FF80FF', '#800080'
])

norm = BoundaryNorm(levels, ncolors=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
)

# The time to visualise
vis_time = basetime + nowcast_interval

for i, mem in enumerate(ltg_frames.ensemble.values):
    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(ensemble=mem, time=vis_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(mem.split("-")[1], fontsize=10)

    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=vis_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, ticks=levels[:-1], extend='max', format='%.3g')
cbar.ax.set_ylabel(ltg_frames.attrs['long_name'], rotation=90)

fig.savefig(
    os.path.join(
        THIS_DIR,
        "../tests/outputs/ltg-hk.png"
    ),
    bbox_inches='tight'
)

###############################################################################
# Deriving lightning probabilities
# ------------------------------------------------------------------------------

# Finally, it is possible to derive
# the probability of lightning exceeding threshold from our ensemble forecast.
#

# Compute exceedance probabililities for a 1 strikes/km^2/h threshold
P = (ltg_frames > 1).sum(dim='ensemble') / len(ltg_frames.ensemble) * 100

# Define cmap, norm
cmap = plt.get_cmap("OrRd", 10)
levels = [
    0, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100
]
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

# Plotting the Exceedance Probability for the first 30 minutes nowcast
fig = plt.figure(figsize=(15, 10))
ax = plt.axes(projection=area_def.to_cartopy_crs())

plot_base(ax, extents, area_def.to_cartopy_crs())

P.where(P > levels[0]).sel(time=vis_time).plot(
    cmap=cmap,
    norm=norm
)

ax.text(
    extents[0],
    extents[3],
    textwrap.dedent(
        """
        Probability exceeding 1 strikes/$km^2$/hr
        Based @ {baseTime}
        """
    ).format(
        baseTime=basetime.strftime('%H:%MH')
    ).strip(),
    fontsize=20,
    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=vis_time.strftime('%H:%MH')
    ).strip(),
    fontsize=20,
    va='bottom',
    ha='right',
    linespacing=1
)

ax.set_title("")

fig.savefig(
    os.path.join(
        THIS_DIR,
        "../tests/outputs/ltg-hk-p.png"
    ),
    bbox_inches='tight'
)