import regionmask

regionmask.__version__

"""
Explanation: Create 2D integer masks
In this tutorial we will show how to create 2D integer mask for arbitrary latitude and longitude grids.
Import regionmask and check the version:
End of explanation
"""


import xarray as xr
import numpy as np

xr.set_options(display_style="text")

"""
Explanation: Load xarray and the tutorial data:
End of explanation
"""


lon = np.arange(-179.5, 180)
lat = np.arange(-89.5, 90)

"""
Explanation: Creating a mask
Define a lon/ lat grid with a 1° grid spacing, where the points define the center of the grid.
End of explanation
"""


regionmask.defined_regions.srex

"""
Explanation: We will create a mask with the SREX regions (Seneviratne et al., 2012).
End of explanation
"""


mask = regionmask.defined_regions.srex.mask(lon, lat)
mask

"""
Explanation: The function mask determines which gridpoints lie within the polygon making up the each region:
End of explanation
"""


import cartopy.crs as ccrs
import matplotlib.pyplot as plt

f, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()))
ax.coastlines()

regionmask.defined_regions.srex.plot(
    ax=ax, add_label=False, line_kws=dict(lw=0.5, color="0.5")
)

mask.plot(ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False)

"""
Explanation: mask is now a xarray.Dataset with shape lat x lon (if you need a numpy array use mask.values). Gridpoints that do not fall in a region are NaN, the gridpoints that fall in a region are encoded with the number of the region (here 1 to 26).
We can now plot the mask:
End of explanation
"""


airtemps = xr.tutorial.load_dataset("air_temperature")

"""
Explanation: Working with a mask
masks can be used to select data in a certain region and to calculate regional averages - let's illustrate this with a 'real' dataset:
End of explanation
"""


# choose a good projection for regional maps
proj = ccrs.LambertConformal(central_longitude=-100)

ax = plt.subplot(111, projection=proj)

airtemps.isel(time=1).air.plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree())

ax.coastlines();

"""
Explanation: The example data is a temperature field over North America. Let's plot the first time step:
End of explanation
"""


mask = regionmask.defined_regions.srex.mask(airtemps)

lon = airtemps.lon.values
print("Grid extent:    {:3.0f}°E to {:3.0f}°E".format(lon.min(), lon.max()))

bounds = regionmask.defined_regions.srex.bounds_global
print("Region extent: {:3.0f}°E to {:3.0f}°E".format(bounds[0], bounds[2]))

"""
Explanation: Conveniently we can directly pass an xarray object to the mask function. It gets the longitude and latitude from the DataArray/ Dataset and creates the mask. Per default regionmask assumes the longitude and latitude are called "lon" and "lat". If they have another name, you can pass them individually, e.g. region.mask(ds.longitude, ds.latitude).
End of explanation
"""


proj = ccrs.LambertConformal(central_longitude=-100)
ax = plt.subplot(111, projection=proj)

low = mask.min()
high = mask.max()

levels = np.arange(low - 0.5, high + 1)

h = mask.plot.pcolormesh(
    ax=ax, transform=ccrs.PlateCarree(), levels=levels, add_colorbar=False
)

# for colorbar: find abbreviations of all regions that were selected
reg = np.unique(mask.values)
reg = reg[~np.isnan(reg)]
abbrevs = regionmask.defined_regions.srex[reg].abbrevs

cbar = plt.colorbar(h, orientation="horizontal", fraction=0.075, pad=0.05)

cbar.set_ticks(reg)
cbar.set_ticklabels(abbrevs)
cbar.set_label("Region")

ax.coastlines()

# fine tune the extent
ax.set_extent([200, 330, 10, 75], crs=ccrs.PlateCarree())

"""
Explanation: Let's plot the mask of the regions:
End of explanation
"""


CNA_index = regionmask.defined_regions.srex.map_keys("C. North America")
CNA_index

"""
Explanation: We want to select the region 'Central North America'. Thus we first need to find out which number this is:
End of explanation
"""


airtemps_CNA = airtemps.where(mask == CNA_index)

"""
Explanation: Mask out a region
xarray provides the handy where method (documentation), that filters elements from this object according to a condition:
End of explanation
"""


# choose a good projection for regional maps
proj = ccrs.LambertConformal(central_longitude=-100)

ax = plt.subplot(111, projection=proj)

regionmask.defined_regions.srex[["CNA"]].plot(ax=ax, add_label=False)

airtemps_CNA.isel(time=1).air.plot.pcolormesh(ax=ax, transform=ccrs.PlateCarree())

ax.coastlines();

"""
Explanation: Check everything went well by repeating the first plot with the selected region:
End of explanation
"""


weights = np.cos(np.deg2rad(airtemps.lat))

ts_airtemps_CNA = airtemps_CNA.weighted(weights).mean(dim=("lat", "lon")) - 273.15

"""
Explanation: Looks good - with this we can calculate the region average.
Calculate weighted regional average
From version 0.15.1 xarray includes a function to calculate the weighted mean - we use cos(lat) as proxy of the grid cell area
End of explanation
"""


f, ax = plt.subplots()
ts_airtemps_CNA.air.plot.line(ax=ax, label="Central North America")

ax.axhline(0, color="0.1", lw=0.5)

plt.legend();

"""
Explanation: We plot the resulting time series:
End of explanation
"""


# you can group over all integer values of the mask
airtemps_all = airtemps.groupby(mask).mean()
airtemps_all

"""
Explanation: To get the regional average for each region you would need to loop over them. However, it's easier to use a 3D mask.
Calculate regional statistics using groupby
End of explanation
"""


# you can group over all integer values of the mask
airtemps_reg_median = airtemps.groupby(mask).median()
airtemps_reg_median.isel(time=0)

"""
Explanation: However, groupby is the way to go when calculating a (unweighted) regional median:
End of explanation
"""


rasm = xr.tutorial.load_dataset("rasm")

"""
Explanation: Multidimensional coordinates
Regionmask can also handle mutltidimensional longitude/ latitude grids (e.g. from a regional climate model). As xarray provides such an example dataset, we will use it to illustrate it. See also in the xarray documentation.
Load the tutorial data:
End of explanation
"""


# choose a projection
proj = ccrs.NorthPolarStereo()

ax = plt.subplot(111, projection=proj)
ax.set_global()

rasm.isel(time=1).Tair.plot.pcolormesh(
    ax=ax, x="xc", y="yc", transform=ccrs.PlateCarree()
)

# add the abbreviation of the regions
regionmask.defined_regions.srex[[1, 2, 11, 12, 18]].plot(
    ax=ax, add_coastlines=False, label="abbrev"
)

ax.set_extent([-180, 180, 43, 90], ccrs.PlateCarree())

ax.coastlines();

"""
Explanation: The example data is a temperature field over the Northern Hemisphere. Let's plot the first time step:
End of explanation
"""


mask = regionmask.defined_regions.srex.mask(rasm.xc, rasm.yc)
mask

"""
Explanation: Again we pass the xarray object to regionmask. We have to pass rasm.xc and rasm.yc as the longitude and latitude coordinates of the array:
End of explanation
"""


rasm_NAS = rasm.where(mask == regionmask.defined_regions.srex.map_keys("NAS"))

"""
Explanation: We want to select the region 'NAS' (Northern Asia).
Select using where
We have to select by index (the number of the region), we thus map from the abbreviation to the index.
End of explanation
"""


# choose a projection
proj = ccrs.NorthPolarStereo()

ax = plt.subplot(111, projection=proj)
ax.set_global()

rasm_NAS.isel(time=1).Tair.plot.pcolormesh(
    ax=ax, x="xc", y="yc", transform=ccrs.PlateCarree()
)


# add the abbreviation of the regions
regionmask.defined_regions.srex[["NAS"]].plot(
    ax=ax, add_coastlines=False, label="abbrev"
)

ax.set_extent([-180, 180, 45, 90], ccrs.PlateCarree())

ax.coastlines();

"""
Explanation: Check everything went well by repeating the first plot with the selected region:
End of explanation
"""


