"""
Radar QPE (Quantitative Precipitation Estimate)
===============================================
This example demonstrates how to generate radar QPE
"""

import os
import xarray as xr
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pyresample import utils

from matplotlib.colors import BoundaryNorm

from swirlspy.rad.iris import read_iris
from swirlspy.qpe.utils import dbz2rr, rr2rf, timestamps_ending, locate_file

plt.switch_backend('agg')

FILE_LOC = os.getcwd()
THIS_DIR = FILE_LOC+'/../tests/int'
os.chdir(THIS_DIR)

############################################################
# Read IRIS RAW data files into xarray
# ---------------------------------------------------
#
# 1. Generate timestamps with timestamps_ending()
# 2. Locate radar data files with locate_file()
# 3. Define target grid as a pyresample AreaDefinition
# 4. Read radar data files with read_iris()
# 5. Convert reflectivity in dBZ to rainrate in mm/h with dbz2rr()
# 6. Convert rainrate to rainfall in 6 mins with rr2rf()
# 7. Concatenate all xarrays into one with xarray.concat()
# 8. Compute the sum along the time dimension

dir = THIS_DIR+"/../samples/iris"
located_files = []
timestamps = timestamps_ending(pd.Timestamp(2018, 8, 22, 19, 00))
for timestamp in timestamps:
    located_files.append(locate_file(dir, timestamp))

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

rainfalls = []
for filename in located_files:
    reflectivity = read_iris(
        filename,
        area_def=area_def,
        coord_label=['easting', 'northing'],
        radar_radius=250000
    )
    rainrate = dbz2rr(reflectivity, a=58.53, b=1.56)
    rainfall = rr2rf(rainrate, time_shift=pd.Timedelta(minutes=6))
    rainfalls.append(rainfall)

rainfall_cat = xr.concat(rainfalls, dim='time')
rainfall_hourly = rainfall_cat.sum(dim='time')

############################################################
# Generate rainfall map
# ---------------------------------------------------
# Define colour scale and format of the rainfall map plot
# then generate it with xarray.plot()

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, 10))
plt.axis('equal')
rainfall_hourly.plot(cmap=cmap, norm=norm)
plt.savefig("radar_qpe.png")