"""
QPF
========================================================
This example demonstrates how to perform
operational deterministic QPF up to three hours from
raingauge and radar data.
"""

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

import os
import numpy as np
import pandas as pd
from pyresample import utils
import xarray
import cartopy.feature as cfeature
import cartopy.crs as ccrs
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm

from swirlspy.rad.irisref import read_irisref
from swirlspy.qpe.utils import dbz2rr, rr2rf, locate_file, timestamps_ending
from swirlspy.qpe.utils import multiple_acc
from swirlspy.obs.rain import Rain
from swirlspy.qpf.rover import rover
from swirlspy.qpf.sla import sla
from swirlspy.core.resample import grid_resample

plt.switch_backend('agg')

THIS_DIR = os.getcwd()
os.chdir(THIS_DIR)

start_time = pd.Timestamp.now()

#############################################################
# Initialising
# ---------------------------------------------------
#
# This section demonstrates extracting
# radar reflectivity data.
#
# Step 1: Define a basetime.
#

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

##############################################################################
# Step 2: Using basetime, generate timestamps of desired radar files
# timestamps_ending() and locate files using locate_file().

# Obtain radar files
dir = THIS_DIR + '/../tests/samples/iris/ppi'
located_files = []
radar_ts = timestamps_ending(
    basetime+pd.Timedelta(minutes=6),
    duration=pd.Timedelta(minutes=60)
)

for timestamp in radar_ts:
    located_files.append(locate_file(dir, timestamp))

##############################################################################
# Step 3: Read data from radar files into xarray.DataArray
# using read_irisref().
#

reflectivity_list = []  # stores reflec from read_iris()
for filename in located_files:
    reflec = read_irisref(
        filename
    )
    reflectivity_list.append(reflec)

########################################################################
# Step 4: Define the target grid as a pyresample AreaDefinition.
# Since the data is in Centered Azimuthal Projection, the source grid
# also must also be defined as a pyresample AreaDefinition.
#

# Defining target 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 = utils.get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

# Defining source grid
radius = 256000
sitecoords = reflectivity_list[0].attrs['site']
res = 480
area_id = 'aeqd'
description = ("Azimuthal Equidistant Projection "
               "centered at the radar site "
               "extending up to {radius:f}m "
               "in each direction "
               "with a {res:f}x{res:f} grid resolution ").format(
                 radius=radius, res=res
                   )
proj_id = 'aeqd'
projection = ('+proj=aeqd +lon_0={lon:f} ' +
              '+lat_0={lat:f} +ellps=WGS84 +datum=WGS84 ' +
              '+units=m +no_defs').format(
                lon=sitecoords[0], lat=sitecoords[1])
x_size = res
y_size = res
area_extent = (-radius, -radius, radius, radius)
area_def_src = utils.get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

##############################################################################
# Step 5: Reproject the radar data from read_irisref() from Centered Azimuthal
# (source) projection to HK 1980 (target) projection.

reproj_reflectivity_list = []
for reflec in reflectivity_list:
    reproj_reflec = grid_resample(
        reflec, area_def_src, area_def_tgt,
        coord_label=['easting', 'northing']
    )
    reproj_reflectivity_list.append(reproj_reflec)

############################################################################
# Step 6: Assigning reflectivity xarrays at the last two timestamps to
# variables for use during ROVER QPF.

xarray1 = reproj_reflectivity_list[-2]
xarray2 = reproj_reflectivity_list[-1]

initialising_time = pd.Timestamp.now()

###########################################################
# Running ROVER and Semi-Lagrangian Advection
# -------------------------------------------
#
# 1. Concatenate two reflectivity xarrays along time dimension.
# 2. Run ROVER, with the concatenated xarray as the input.
# 3. Perform Semi-Lagrangian Advection using the motion fields from rover.

# Combining the two reflectivity DataArrays
# the order of the coordinate keys is now ['y', 'x', 'time']
# as opposed to ['time', 'x', 'y']
reflec_concat = xarray.concat([xarray1, xarray2], dim='time')

# Rover
motion_u, motion_v = rover(
    reflec_concat
)

rover_time = pd.Timestamp.now()

# Semi Lagrangian Advection
reflectivity = sla(
    reflec_concat, motion_u, motion_v, steps=30
)

sla_time = pd.Timestamp.now()

##############################################################
# Concatenating observed and forecasted reflectivities
# ---------------------------------------------------
#
# 1. Add forecasted reflectivity to reproj_reflectivity_list.
# 2. Concatenate observed and forecasted reflectivity
#    xarray.DataArrays along the time dimension.

reproj_reflectivity_list.append(reflectivity[1:, ...])
reflectivity = xarray.concat(reproj_reflectivity_list, dim='time')
reflectivity.attrs['long_name'] = 'reflectivity'

concat_time = pd.Timestamp.now()

###################################################
# Generating labelled radar colour images
# ----------------------------------------
#
# Define the colour scale and the format of the
# reflectivity plots. Then generate using xarray.plot().
# In this example, only hourly images will be plotted.

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

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

# generate timesteps for plotting
timesteps = [basetime+pd.Timedelta(hours=i) for i in range(4)]

# Plot radar colour images. Only plot hourly images
for time in timesteps:
    plt.figure(figsize=(15, 10))
    plt.axis("equal")
    reflectivity.sel(time=time).plot(cmap=cmap, norm=norm)
    plt.savefig(
        os.path.join(
            THIS_DIR,
            ("../tests/outputs/rover-output-"
             f"{time.strftime('%Y%m%d%H%M')}.png")
            )
        )

#############################################
# Generating radar reflectivity maps
# -----------------------------------
#
# Using the colour scale generated in the previous section,
# generate radar reflectivity maps with Natural Earth Feature Coastlines
# and motion field.
# In this example only hourly images will be plotted.

# Defining quiver parameters
qx = motion_u.coords['easting'].values[::5]
qy = motion_u.coords['northing'].values[::-5]
qu = motion_u.values[::5, ::5]
qv = motion_v.values[::5, ::5]

for time in timesteps:
    plt.figure(figsize=(15, 10))
    plt.axis("equal")
    crs = area_def_tgt.to_cartopy_crs()
    ax = plt.axes(projection=crs)
    ax.coastlines(resolution='10m')
    ax.gridlines()
    reflectivity.sel(time=time).plot(cmap=cmap, norm=norm)
    ax.quiver(
        qx, qy, qu, qv, pivot='mid', regrid_shape=20
    )
    plt.savefig(
        os.path.join(
            THIS_DIR,
            ("../tests/outputs/rover-output-map-"
             f"{time.strftime('%Y%m%d%H%M')}.png")
        )
    )

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+0min.
#
# | Step 1: Convert reflectivity in dBZ to rainrates in mm/h with dbz2rr().
# | Step 2: Convert rainrates to rainfalls in 6 mins with rr2rf().

rainrates = dbz2rr(reflectivity, a=58.53, b=1.56)
rainfalls = rr2rf(rainrates)

######################################################################
# Step 3: Using multiple_acc(), calculate hourly accumulated
# rainfall every 30 minutes.

acc_rf = multiple_acc(
    rainfalls, basetime, basetime+pd.Timedelta(hours=3)
)

acc_time = pd.Timestamp.now()

#######################################################################
# Plotting rainfall maps
# ---------------------------------------
#
# 1. Define the colour scheme.
# 2. Defining the projection.
# 3. Defining the zoom or area extent. Tuple order is (x0, x1, y0, y1)
#    as opposed to pyresample (x0, y0, x1, y1).
# 4. Define figure
# 5. Define axes.
# 6. Adding area extent defined in (3) to axes.
# 7. Adding coastlines (this example uses GSHHS).
# 8. Add gridlines.
# 9. Plot xarray using xarray.plot().
# 10. Adding title (if required).

# Defining the colour scheme
levels = [
    0, 0.5, 2, 5, 10, 20,
    30, 40, 50, 70, 100, 150,
    200, 300, 400, 500, 600, 700
]

cmap = matplotlib.colors.ListedColormap([
    '#ffffff', '#9bf7f7', '#00ffff', '#00d5cc', '#00bd3d', '#2fd646',
    '#9de843', '#ffdd41', '#ffac33', '#ff621e', '#d23211', '#9d0063',
    '#e300ae', '#ff00ce', '#ff57da', '#ff8de6', '#ffe4fd'
])

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

# Defining projection
crs = area_def_tgt.to_cartopy_crs()
# Defining zoom extent
r = 64000
proj_site = xarray1.proj_site
zoom = (
    proj_site[0]-r, proj_site[0]+r, proj_site[1]-r, proj_site[1]+r
)
# Defining coastlines
hires = cfeature.GSHHSFeature(
            scale='h',
            levels=[1],
            edgecolor='black',
            facecolor='none'
        )

for time in acc_rf.coords['time'].values:
    # Define figure
    plt.figure(figsize=(15, 10))
    plt.axis('equal')
    # Plotting axes
    ax = plt.axes(projection=crs)
    # setting extent
    ax.set_extent(zoom, crs=crs)
    # Adding coastlines
    ax.add_feature(hires)
    # Add gridlines
    ax.gridlines()

    # Plot data
    acc_rf.sel(time=time).plot(cmap=cmap, norm=norm)

    # Add title
    t_minus = pd.Timestamp(time) - \
        basetime-pd.Timedelta(minutes=60)
    t_minus = t_minus.to_timedelta64().astype('timedelta64[m]')
    t_plus = pd.Timestamp(time) - basetime
    t_plus = t_plus.to_timedelta64().astype('timedelta64[m]')
    plt.title(
        "Quantitative Precipitation Forecast with Hourly"
        " Accumulated Rainfall.\nBasetime: " +
        str(basetime) + "\nStart time: t " +
        str(t_minus) + " End time: t " +
        str(t_plus)
    )
    plt.savefig(
        THIS_DIR +
        "/../tests/outputs/forecast-" +
        f"{pd.Timestamp(time).strftime('%Y%m%d%H%M')}.png"
        )

rf_image_time = pd.Timestamp.now()

##############################################################################
# Extract the rainfall values at a specified location
# ------------------------------------------------------------------
#
# In this example, the rainfall values at the location is assumed
# to be the same as the nearest gridpoint.
#
# 1. Read information regarding the radar stations into a pandas.DataFrame.
# 2. Extract the rainfall values at the nearest gridpoint to location
#    for given timesteps (in this example, 30 minute intervals).
# 3. Store rainfall values over time in a pandas.DataFrame.
# 4. Plot the time series of rainfall at different stations.

# Getting radar station coordinates
df = pd.read_csv(
    os.path.join(THIS_DIR, "../tests/samples/hk_raingauge.csv"),
    usecols=[0, 1, 2, 3, 4]
)

# Extract rainfall values at gridpoint closest to the
# location specified for given timesteps and storing it
# in pandas.DataFrame.

rf_time = []
for time in acc_rf.coords['time'].values:
    rf = []
    for index, row in df.iterrows():
        rf.append(acc_rf.sel(
            time=time, northing=row[1],
            easting=row[2],
            method='nearest'
        ).values)
    rf_time.append(rf)

rf_time = np.array(rf_time)

station_rf = pd.DataFrame(
    data=rf_time,
    columns=df.iloc[:, 0],
    index=pd.Index(
        acc_rf.coords['time'].values,
        name='time'
    )
)
print(station_rf)

# Plotting time series graph
ax = station_rf.plot(title="Time Series of Hourly Accumulated Rainfall",
                     grid=True)
ax.set_ylabel("Hourly Accumulated Rainfall (mm)")
plt.savefig(THIS_DIR+"/../tests/outputs/qpf_time_series.png")

extract_time = pd.Timestamp.now()

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

print(f"Start time: {start_time}")
print(f"Initialising time: {initialising_time}")
print(f"Rover time: {rover_time}")
print(f"SLA time: {sla_time}")
print(f"Concatenating time: {concat_time}")
print(f"Plotting radar image time: {radar_image_time}")
print(f"Accumulating rainfall time: {acc_time}")
print(f"Plotting rainfall map time: {rf_image_time}")
print(f"Extracting and plotting time series time: {extract_time}")

print(f"Time to initialise: {initialising_time-start_time}")
print(f"Time to run rover: {rover_time-initialising_time}")
print(f"Time to perform SLA: {sla_time-rover_time}")
print(f"Time to concatenate xarrays: {concat_time - sla_time}")
print(f"Time to plot radar image: {radar_image_time - concat_time}")
print(f"Time to accumulate rainfall: {acc_time - radar_image_time}")
print(f"Time to plot rainfall maps: {rf_image_time-acc_time}")
print(f"Time to extract and plot time series: {extract_time-rf_image_time}")