"""
Quantitative Precipitation Estimate (QPE)
====================================================

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

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


import swirlspy.qpe.rfmap as rfmap
from swirlspy.obs.rain import Rain
from swirlspy.qpe.utils import locate_file, timestamps_ending

THIS_DIR = os.getcwd()

###############################################################################
# Initialising
# ---------------------------------------
#
# 1. Generate timestamps with timestamps_ending()
# 2. Locate rainfall data files with locate_file()
# 3. Define target grid as a pyresample AreaDefinition
# 4. Construct a Rain object containing rainfall data


dir = THIS_DIR+"/../tests/samples"
located_files = []
timestamps = timestamps_ending(
    pd.Timestamp(2018, 8, 11, 00, 40),
    duration=pd.Timedelta('1 hour'),
    interval=pd.Timedelta('10 minutes'),
    exclude_end=False
    )

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

# Writing AreaDefinition of target grid
area_id = "WGS84"
description = "Plate Carree"
proj_id = "WGS84"
projection = '+proj=longlat +ellps=WGS84 +datum=WGS84 +units=degrees'
x_size = 500
y_size = 500
area_extent = (
    119.8, 13.6, 122.0, 16.5
    )
area_def = utils.get_area_def(
    area_id, description, proj_id, projection, x_size, y_size, area_extent
)

# Constructing a Rain object using data from the text file
rg_object = Rain(
    located_files,
    proj="Plate Carree",
    duration=pd.Timedelta('10 minutes'),
    NAN=3276.7
)


###############################################################################
# Interpolation
# -------------
#
# Perform RBF interpolation. Interpolated rainfall is returned
# along with metadata as an xarray.DataArray.

# Interpolate into target grid specified by AreaDefinition
interpolated_rf = rfmap.rg_interpolate(
    rg_object, area_def, 'rbf',
    coord_label=['Longitude', 'Latitude']
)
###############################################################################
# Generate rainfall map
# ---------------------
#
# Define the colour scale and format of the reflectivity plots.
# Generate the image using xarray.DataArray.plot().

# Plotting
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)
plt.figure(figsize=(15, 15))
plt.axis("equal")

# AreaDefinition.to_cartopy_crs() does not work for
# Geographic Coordinate Systems
ax = plt.axes(projection=ccrs.PlateCarree())
ax.set_extent(
    [area_extent[0], area_extent[2],
        area_extent[1], area_extent[3]]
)
ax.coastlines(resolution='10m')  # plot coastlines
# Plot data
interpolated_rf.plot(
    cmap=cmap, norm=norm, extend='neither'
    )
# Show stations
rfmap.show_raingauge(rg_object)
# Plot title
plt.title(
    "Accumulated rainfall from "
    f"{rg_object.start_time} to {rg_object.end_time}"
)
# timestring
start_timestr = rg_object.start_time.strftime("%Y%m%d%H%M")
end_timestr = rg_object.end_time.strftime("%Y%m%d%H%M")
plt.savefig(
    os.path.join(
        THIS_DIR,
        "../tests/outputs",
        f"RBF_{start_timestr}_{end_timestr}.png"
    )
    )