"""
QPE (Manila)
====================================================

This example demonstrates how to perform QPE,
using raingauge data from Manila and radar data from Subic.
"""

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

import os
import time
import numpy as np
import pandas as pd
import copy
import xarray
import scipy
import pyproj
import matplotlib
from matplotlib.colors import BoundaryNorm
import matplotlib.pyplot as plt
from PIL import Image
from pyresample import utils
from cartopy.io import shapereader
import cartopy.feature as cfeature
import cartopy.crs as ccrs
import cartopy.io.shapereader as shpreader

from swirlspy.obs import Rain
from swirlspy.rad.uf_ph import read_uf_ph
from swirlspy.core.resample import grid_resample
from swirlspy.qpe.rfmap import rg_interpolate, comp_qpe, show_raingauge
from swirlspy.qpe.utils import timestamps_ending, locate_file, dbz2rr, rr2rf, \
    temporal_interp, acc

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


###################################################################
# Initialising
# -----------------------------------------------------------------
#
# This section demonstrates extracting raingauge and radar data.
#
# Step 1: Defining an end-time for accumulating rainfall.
#

acctime = pd.Timestamp('20180811112000').floor('min')
acctime_str = acctime.strftime('%Y%m%d%H%M')

####################################################################
# Step 2: Setting up directories for raingauge and radar files.
#

rad_dir = THIS_DIR + '/../tests/samples/uf_ph/sub/'
rg_dir = THIS_DIR + '/../tests/samples/rfmap/'

###################################################################
# Step 3: Generating timestamps and pattern for both radar and
# raingauge files.

# Timestamps of raingauges
rg_timestrings = timestamps_ending(
    acctime + pd.Timedelta(minutes=10),
    duration=pd.Timedelta(hours=1),
    interval=pd.Timedelta(minutes=10)
)

# Raingauge pattern
rg_pattern = ['rf60m_20'+ts for ts in rg_timestrings]

# Finding time nearest radar file
# to accumulation end time
minute = acctime.minute
nearest_6min = acctime.minute // 10 * 10

nearest_rad_timestamp = pd.Timestamp(
    acctime_str[:-2]+f'{nearest_6min:02}'
)

rad_timestrings = timestamps_ending(
    nearest_rad_timestamp,
    duration=pd.Timedelta(hours=1),
    interval=pd.Timedelta(minutes=10)
)

#####################################################################
# Step 4: Extracting raingauge and radar files from
# their respective directories.

located_rg_files = []
for pat in rg_pattern:
    located_rg_files.append(locate_file(rg_dir, pat))


located_radar_files = []
for ts in rad_timestrings:
    located_radar_files.append(locate_file(rad_dir, ts))


#####################################################################
# Step 5: Read data from raingauge files into a Rain object.
# Coordinates are geodetic, following that in the files.

rg_object_geodetic = Rain(
    located_rg_files,
    'WGS84',
    duration=pd.Timedelta(minutes=5),
    NAN=[3276.7, 32767]
)

######################################################################
# Step 6: Define the target grid as a pyresample AreaDefinition.

# Defining target grid
area_id = "epsg3123_240km"
description = ("A 240 m resolution rectangular grid "
               "centred at Subic RADAR and extending to 240 km "
               "in each direction")
proj_id = 'epsg3123'
projection = ('+proj=tmerc +lat_0=0 '
              '+lon_0=121 +k=0.99995 +x_0=500000 '
              '+y_0=0 +ellps=clrk66 +towgs84=-127.62,-67.24,'
              '-47.04,-3.068,4.903,1.578,-1.06 +units=m '
              '+no_defs')
x_size = 1000
y_size = 1000
area_extent = (191376.04113, 1399386.68659, 671376.04113, 1879386.68659)
area_def = utils.get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

#########################################################################
# Step 7: Convert coordinates of raingauge object to desired projection.
# In this example, the desired projection is PRS92.
# This can be achieved by the .reproject() method of the Rain object.
#

inProj = pyproj.Proj(init="epsg:4326")
outProj = pyproj.Proj(area_def.proj_str)

rg_object = rg_object_geodetic.reproject(inProj, outProj, "PRS92")

######################################################################
# Step 8: Read radar files into xarray.DataArrays using read_uf_ph().

reflec_list = []
for file in located_radar_files:
    reflec = read_uf_ph(
        file, area_def=area_def,
        coord_label=['easting', 'northing'],
        indicator='deg_ppi', elevation=0.5
    )

    reflec_list.append(reflec)

reflectivity = xarray.concat(reflec_list, dim='time')

###################################################################
# Accumulating and interpolating rainfall
# -----------------------------------------------------------------
#
# Interpolate rainfall recorded by raingauges into the user-defined grid
# and accumulate radar rainfall over an hour
# after making the necessary
# adjustments.
#

##################################################################
# Step 1: Interpolate Rain object to user-defined grid.
# In this example, a multiquadric Radial Basis Function
# is used.
#

interpolated_rg = rg_interpolate(
    rg_object, area_def, 'rbf',
    coord_label=['easting', 'northing']
    )


##############################################################################
# Step 2: Convert to radar reflectivity to rainrates,
# convert rainrates to times of raingauges,
# and accumulate rainfalls every 10 minutes.

rainrates = dbz2rr(reflectivity, a=300, b=1.4)
# Convert time coordinates of rainrates from start time
# to end time
rainrates_time_endtime = rainrates.copy()
rainrates_time_endtime.coords['time'] = \
    [
        pd.Timestamp(t) + pd.Timedelta(minutes=10)
        for t in rainrates.coords['time'].values
    ]

rainfalls = rr2rf(rainrates_time_endtime, scan_duration=10)

###############################################################################
# Step 3: Accumulate rainfall over an hour.

acc_rf = acc(
    rainfalls,
    rg_object.end_time,
    acc_period=pd.Timedelta(minutes=60)
)

###################################################################
# Compositing rainfall
# -----------------------------------------------------------------
#
# Perform compositing on radar and raingauge derived rainfall
# to obtain a composite QPE.

comprf = comp_qpe(
    area_def,
    rg_object=rg_object,
    rg_interp=interpolated_rg,
    rad_rf=acc_rf
)

###################################################################
# Plotting
# ---------------------------------------------------------------
#
# Plot composited radar and raingauge rainfall.
#

# Defining coastlines
map_shape_file = os.path.join(THIS_DIR, "./../tests/samples/shape/rsmc")
ocean_color = np.array([[[178, 208, 254]]], dtype=np.uint8)
land_color = cfeature.COLORS['land']
coastline = cfeature.ShapelyFeature(
    list(shpreader.Reader(map_shape_file).geometries()),
    ccrs.PlateCarree()
)

# Plotting function for neatness.

def plot_map(
    da, rg_object, acctime, area_def,
    based='raingauge and radar',
    savepath='', area_extent=None
):

    """
    A custom function for plotting a map.

    Parameters
    --------------
    da: xarray.DataArray
        Contains data to be plotted.
    rg_object: Rain
        Contains raingauge data.
    acctime: pd.Timestamp
        Contains the endtime of the accumulation
        period.
    area_def: pyresample.geometry.AreaDefinition
        AreaDefinition of the grid.
    based: str
        Type of data plotted in the map.
    savepath: str
        Path to save the image to.
    area_extent: tuple
        Area extent (x0, x1, y0, y1) to be plotted.
        Defaults to None.

    """
    # 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)

    # Plotting axes
    plt.figure(figsize=(24, 21))
    crs = area_def.to_cartopy_crs()
    ax = plt.axes(projection=crs)

    if area_extent is not None:
        ax.set_extent(area_extent, crs=crs)

    # Plot data
    quadmesh = da.plot(
        cmap=cmap,
        norm=norm,
        extend='max',
        cbar_kwargs={'ticks': levels, 'format': '%.3g'}
        )

    # Adjusting size of colorbar
    cb = quadmesh.colorbar
    cb.ax.set_ylabel(
        da.attrs['long_name']+'['+da.attrs['units']+']',
        fontsize=28
    )

    cb.ax.tick_params(labelsize=24)

    # Setting labels
    ax.xaxis.set_visible(True)
    ax.yaxis.set_visible(True)

    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(24)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(24)

    ax.xaxis.label.set_size(28)
    ax.yaxis.label.set_size(28)

    # ocean
    ax.imshow(np.tile(ocean_color, [2, 2, 1]),
            origin='upper',
            transform=ccrs.PlateCarree(),
            extent=[-180, 180, -180, 180],
            zorder=-1)
    # coastline, color
    ax.add_feature(coastline,
                facecolor=land_color, edgecolor='none', zorder=0)
    # overlay coastline without color
    ax.add_feature(coastline, facecolor='none',
                edgecolor='gray', linewidth=0.5, zorder=3)

    # Show raingauge
    show_raingauge(
        rg_object, ax, show_value=True, color='red', markersize=20,
        fontsize=20
        )
    # Show title
    plt.title(
        (f"Last Hour Rainfall at {acctime.strftime('%H:%MH %d-%b-%Y')}\n"
         f"(based on {based} data)"),
        fontsize=32
    )

    plt.savefig(savepath)

############################################################################

# Plotting maps

r = 64000
proj_site = reflectivity.proj_site
area_extent = (
    proj_site[0]-r,
    proj_site[0]+r,
    proj_site[1]-r,
    proj_site[1]+r
)

# Raingauge only
plot_map(
    interpolated_rg, rg_object,
    acctime, area_def,
    based='raingauge',
    savepath=THIS_DIR+f'/../tests/outputs/raingauge_{acctime_str}.png',
    area_extent=area_extent
)

# Radar only
plot_map(
    acc_rf, rg_object,
    acctime, area_def,
    based='radar',
    savepath=THIS_DIR+f'/../tests/outputs/radar_{acctime_str}.png',
    area_extent=area_extent
)

# Composite raingauge and radar
plot_map(
    comprf, rg_object,
    acctime, area_def,
    savepath=THIS_DIR+f'/../tests/outputs/comp_{acctime_str}.png',
    area_extent=area_extent
)