# Figure 7 of paper "A locally smoothed terrain-following coordinate"
# submitted to JAMES
# Author: Stephanie Westerhuis
# Date: February 1, 2021

import matplotlib
import numpy as np
import xarray as xr
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import datetime as dt
import sys
from scipy import interpolate, stats
import math
from pathlib import Path
import cartopy
import matplotlib.image as mpimg
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import seaborn as sns
import matplotlib.path as mpath

# illustrate satellite-based verification algorithm
# example case: observed FLS fraction on December 29, 2016, 09 UTC

def add_simple_background(fig, region='switzerland'):

    borders = cartopy.feature.NaturalEarthFeature(
        category='cultural',
        name='admin_0_countries_lakes',
        scale='10m',
        edgecolor='black',
        facecolor='none')

    extent = {
        # lon_left, lon_right, lat_bottom, lat_top
        'switzerland': [5.65, 11, 45.6, 48],
        'switzerland_tight': [5.8, 10.6, 45.75, 47.85],
         'rheintal'  : [6, 11.5, 47, 49.5],
         'zrh'       : [8, 9.2, 47.3, 47.9],
         'mittelland': [5.9, 10.7, 45.9, 48]
         }

    ax = fig.add_subplot(111, projection=cartopy.crs.Mercator())
    ax.add_feature(borders,                    zorder = 1)
    # map extent for Switzerland (lonleft, lonright, latdown, latup)
    ax.set_extent(extent[region]) 
    return ax

def get_mittelland_mask(lats, lons):
    """
    create a boolean mask for the
    swiss mittelland
    input:
    lats, lons: 2D arrays
    output:
    boolean array
    """

    ll_corner = (46.12,5.89)
    p1        = (46.06,6.10)
    p2        = (46.33,6.78)
    p3        = (46.55,7.01)
    p4        = (46.64,7.31)
    p5        = (46.65,7.65)
    p6        = (46.62,7.79)
    p7        = (46.81,8.33)
    p8        = (47.09,9.78)
    p9        = (47.82,10.02)
    p10       = (47.81,8.34)
    p11       = (47.38,7.87)
    p12       = (47.29,7.68)
    p13       = (47.25,7.45)
    p14       = (47.13,7.06)
    p15       = (47.07,6.87)
    p16       = (46.73,6.36)
    p17       = (46.59,6.30)
    p18       = (46.18,5.86)

    # create polygon
    Path = mpath.Path
    path_data = [
        (Path.MOVETO, ll_corner),
        (Path.LINETO, p1),
        (Path.LINETO, p2),
        (Path.LINETO, p3),
        (Path.LINETO, p4),
        (Path.LINETO, p5),
        (Path.LINETO, p6),
        (Path.LINETO, p7),
        (Path.LINETO, p8),
        (Path.LINETO, p9),
        (Path.LINETO, p10),
        (Path.LINETO, p11),
        (Path.LINETO, p12),
        (Path.LINETO, p13),
        (Path.LINETO, p14),
        (Path.LINETO, p15),
        (Path.LINETO, p16),
        (Path.LINETO, p17),
        (Path.LINETO, p18),
        (Path.CLOSEPOLY, ll_corner),
        ]
    codes, verts = zip(*path_data)
    path = mpath.Path(verts, codes)

    # store original shape
    shape = lats.shape

    # path.contains_points checks whether points are within polygon
    # however, this function can only handle vectors
    # -> ravel and unravel
    latlon = [[lat, lon] for lat, lon in zip(lats.ravel(), lons.ravel())]
    mask = np.reshape(path.contains_points(latlon),shape)

    return mask

f = 'data/201612290900_cosmo1_ccl.nc'
ds = xr.open_dataset(f)

fig = plt.figure(figsize=(6,4), dpi=300)
fig.subplots_adjust(left=0.03, bottom=0.03, right=0.97, top=0.97)
ax = add_simple_background(fig, region='switzerland_tight')

lats = ds.latitude.values
lons = ds.longitude.values

# region comprising swiss plateau ("mittelland")
mask_ml_obs = get_mittelland_mask(lats,lons)
values = ds.cloud_confidence_level.where(mask_ml_obs).values

# cloud confidence level > 0.7 -> FLS
ccl_threshold = 0.7
color_levels = [0, ccl_threshold, 1]
ax.contourf(lons, lats, values,
            zorder = 1,
            transform = cartopy.crs.PlateCarree() ,
            levels = color_levels,
            colors = ('silver', 'green'))

# legend
legend_elements = [Patch(facecolor='green', edgecolor='k',
                         label='FLS'),
                   Patch(facecolor='silver', edgecolor='k',
                         label='no FLS')]

ax.legend(handles = legend_elements,
          loc = 'lower right')

# save
out_name = '7_ccl'
fig.savefig(f'plots/{out_name}.png')