#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 11 16:54:46 2021
@author: mika
"""
#mask creation:
import numpy as np
import shapely as shp #no canonical shortname
import shapely.vectorized #vectorized version has to be loaded in separately
import pandas as pd #optional, only for loading in coordinates from csv
#netcdf file creation:
import netCDF4 as nc4
#plotting:
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.patches as mpatches #for artificial legends
import matplotlib.lines as mlines #for artificial legends
#unsure:
#from cartopy import config
#import copy #from python, but because of cmap copying
#if warning due to cartopy and matplotlib versions clashing:
#whether you have this problem and need this will depend on your specific cartopy and matplotlib versions
from matplotlib.axes import Axes
from cartopy.mpl.geoaxes import GeoAxes
GeoAxes._pcolormesh_patched = Axes.pcolormesh
# %% setup global grid and polygon
#global grid, points to check for:
res=0.25 # resolution in degrees
global_lats=np.arange(-90,90+res,res)
global_lons=np.arange(-180,180+res,res)
x, y = np.meshgrid(global_lons,global_lats) #creates meshgrid of size lats*lons
# the polygon, our mask:
#df=pd.read_csv('coords_scandinavia.csv')
#coordinates=list(zip(df.lon,df.lat))
#if not loading in a csv-file. coordinates needs to be a python LIST of TUPLE pairs of (lon,lat) (x,y)
coordinates=\
[(31.25, 69.75), (23.75, 66.0), (20.0, 62.0), (20.0, 55.25), (15.0, 54.75),(13.0, 55.0),
(12.0, 54.5), (10.75, 54.0), (8.75, 54.0), (8.5, 54.4), (4.0, 54.5), (4.0, 72.0), (31.0, 72.0)]
nodes=coordinates.copy() #needs to be copy or otherwise appends updates with the appending below
nodes.append(coordinates[0]) #polygon needs to be closed, so first point is endpointt
#create polygon:
geom=shp.geometry.Polygon(nodes)
# %% calcuate boolean mask:
mask_inner=shapely.vectorized.contains(geom, x, y) # - for interior points
mask_edge=shapely.vectorized.touches(geom, x, y) #- for edge points
mask_full=mask_inner|mask_edge # | is boolean OR bitwise (elementwise) operator. if either of the elements are true, result is true. if both are false, false.
# %% super quick plotting for checking
#matplotlibs imshow is by far the quickest for such dense plotting (1M points if res=0.25
#quicklook plot, NOT in geographical projection
#imshow uses ij indexing, so we use the origin argument
fig, ax = plt.subplots(3, 1, figsize=(14,25))
cmap='bwr'
ax[0].imshow(mask_inner,origin='lower',cmap=cmap)
ax[1].imshow(mask_edge,origin='lower',cmap=cmap)
ax[2].imshow(mask_full,origin='lower',cmap=cmap)
ax[0].set_title('inner mask')
ax[1].set_title('edge mask')
ax[2].set_title('combined mask')
# for ax in axs.ravel():
# ax.grid(b=True,which='major',axis='x',color='k',alpha=0.4,linestyle='--')
# ax.set_xticks(months)
# ax.set_xticklabels([x[0:3] for x in monthnames],rotation=90,ha='center')
# %% geographical global map:
dataproj=ccrs.PlateCarree()
mapproj=ccrs.PlateCarree() #regular, fine for global
fig, ax = plt.subplots(1, 1, subplot_kw=dict(projection=mapproj), figsize=(14,14))
ax.set_global()
#ax.set_extent((3,27,53,73),crs=dataproj) #fine for scandinavia
ax.coastlines('10m') #choices 110m, 50m and 10m. m does not mean meter, but 1:Xmillion.
ax.add_feature(cfeature.BORDERS, linestyle=':')
ax.gridlines(color='k',alpha=0.3,linestyle=':')
pixelplot=ax.pcolormesh(global_lons,global_lats,mask_full,transform=dataproj,cmap='bwr')
ax.set_title('global map proj \n full mask')
# %% zoom-in geographical map standard projection
#we find coordinates to zoom by based on polygon bounds:
minlon,maxlon,minlat,maxlat=geom.bounds
zoomin=geom.bounds
pad=2 #padding in degrees, adjust as needed
dataproj=ccrs.PlateCarree()
mapproj=ccrs.PlateCarree() #regular, fine for global
fig, ax = plt.subplots(1, 1, subplot_kw=dict(projection=mapproj), figsize=(14,14))
#ax.set_global()
ax.set_extent((minlon-pad, maxlon+pad, minlat-pad, maxlat+pad),crs=dataproj)
ax.coastlines('10m') #choices 110m, 50m and 10m. m does not mean meter, but 1:Xmillion.
ax.add_feature(cfeature.BORDERS, linestyle=':')
ax.gridlines(color='k',alpha=0.3,linestyle=':')
pixelplot=ax.pcolormesh(global_lons,global_lats,mask_full,transform=dataproj,cmap='bwr')
ax.set_title('zoom-in standard map proj \n full mask')
#as you will see, this is not very good unless global mask or close to equator
# %% closeup-geograpical, better projection choice
#when switching to LambertConformal projection, or others, it can no longer handle a input data set that is global.
#it will simply plot nothing. Known bug/undesired behavoir in cartopy from cyclic points and transformations, or something like that
#we will therefore make a subset of the data based
minlon,maxlon,minlat,maxlat=geom.bounds
zoomin=geom.bounds
pad=5 #padding in degrees
#xs and ys to subset data with
xss1=np.where(global_lons==(minlon-pad))[0][0]
xss2=np.where(global_lons==(maxlon+pad))[0][0]
yss1=np.where(global_lats==(minlat-pad))[0][0]
yss2=np.where(global_lats==(maxlat+pad))[0][0]
mask_inner_subset=mask_inner[yss1:yss2+1,xss1:xss2+1]
mask_edge_subset=mask_edge[yss1:yss2+1,xss1:xss2+1]
dataproj=ccrs.PlateCarree()
#mapproj=ccrs.PlateCarree()
#this specific mapproj and settings are nice for Scandinavia, specifically.
#test different projections/settings to fit your mask
mapproj=ccrs.LambertConformal(central_longitude=8, central_latitude=60,
false_easting=0.0, false_northing=0.0, secant_latitudes=None,
standard_parallels=None, globe=None, cutoff=-30)
fig, ax = plt.subplots(1, 1, subplot_kw=dict(projection=mapproj), figsize=(14,14))
#ax1.set_global()
ax.set_extent((3,27,53,73),crs=dataproj) #fine for scandinavia
ax.coastlines('10m') #choices 110m, 50m and 10m. m does not mean meter, but 1:Xmillion.
ax.add_feature(cfeature.BORDERS, linestyle=':')
ax.gridlines(color='k',alpha=0.3,linestyle=':')
#matplotlib.colors.to_rgba(c, alpha=None)
alpha=0.3
edgeplot=ax.pcolormesh(global_lons[xss1:xss2+1],global_lats[yss1:yss2+1],mask_inner_subset,
transform=dataproj,cmap='Reds',shading='auto',edgecolor='lightgrey',alpha=alpha,label='_inner')
innerplot=ax.pcolormesh(global_lons[xss1:xss2+1],global_lats[yss1:yss2+1],mask_edge_subset,
transform=dataproj,cmap='Blues',shading='auto',edgecolor='lightgrey',alpha=alpha,label='_edge')
xs=[x[0] for x in coordinates]
ys=[y[1] for y in coordinates]
nodesplot=ax.scatter(xs,ys,s=50,c='yellow',transform=dataproj,marker='o',label='_nodes')
pathplot=ax.add_geometries([geom],dataproj,facecolor='none', edgecolor='green',label='_path')
ax.set_title('better projection zoom in \n inner, edge, nodes and path')
# #make artifical legend elements:
# note: the hidden _labels in the plots are because else it makes plotting really, really show as it checks for a label first
red_patch = mpatches.Patch(color=mcolors.to_rgba('red',alpha=alpha), label='inner points')
blue_patch = mpatches.Patch(color=mcolors.to_rgba('blue',alpha=alpha), label='edge points')
yellow_dots = mlines.Line2D([],[],color='yellow',marker='o',linestyle=None, label='nodes')
black_line = mlines.Line2D([],[],color='green',marker=None,linestyle='-', label='path')
ax.legend(handles=[red_patch,blue_patch,yellow_dots,black_line],loc='upper right')
# %% make netcdf file
#outcommented now. only comment in when making.
#remember you cannot change in an already created netcdf file, so make sure you create a new one first
# homemade_mask=nc4.Dataset('homemade_mask.nc','w',format='NETCDF4')
# lon = homemade_mask.createDimension("lon", global_lons.size)
# lon = homemade_mask.createVariable("lon", "f8",("lon"), zlib=True)
# lon[:]=global_lons #put data in
# lon.units="degrees_east"
# lon.standard_name="longitude"
# lon.long_name="longitude"
# lat = homemade_mask.createDimension("lat", global_lats.size)
# lat = homemade_mask.createVariable("lat", "f8",("lat"), zlib=True)
# lat[:]=global_lats #put data in
# lat.units="degrees_north"
# lat.standard_name="laitude"
# lat.long_name="laitude"
# time = homemade_mask.createDimension("time",None)
# time = homemade_mask.createVariable("time", "f8", ("time"), zlib=True)
# tunits='days since 0001-01-01'
# calendar='standard'
# time.units=tunits
# time.calendar=calendar
# time.standard_name="time"
# time[:]=0 #it needs atleast one time to work with watersip
# mask = homemade_mask.createVariable("mask", "f8", ("time","lat","lon"), zlib=True)
# mask[0,:,:]=mask_full #put the mask data in
# mask.units="mm/day"
# mask.xmin=-180
# mask.xmax=180
# mask.xstag=180
# mask.ymin=-90
# mask.ymax=90
# mask.ystag=90
# # time = mob_flux.createDimension("time",None)
# # times = mob_flux.createVariable("time","f8",("time"),zlib=True)
# # times.units = tunits
# # times.calendar = calendar
# # dates=fluxdata_5s.index.to_pydatetime()
# # times[:] = nc4.date2num(dates, units=tunits, calendar=calendar)
# homemade_mask.close() #remember to close the file
# %% OUTSIDE PYTHON
"""
in a terminal, navigate to where you newly created mask file is
then use
"cdo -f nc2 copy mask-from-python.nc mask-to-watersip.nc
"
"""
