######################################################################
# Import the needed modules.
#
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER


######################################################################
# The following code reads the example data using the xarray open_dataset
# function and prints the coordinate values that are associated with the
# various variables contained within the file.

ds = xr.open_dataset('wrfout_d1.uvrewrite.nc')
ds.coords

# Grab lat/lon values
lats = ds.XLAT.data[0]
lons = ds.XLONG.data[0]
lllon = lons[0][0]
lllat = lats[0][0]
urlon = lons[-1][-1]
urlat = lats[-1][-1]

# Select and grab data
t2m  = ds['T2'].data[0]
uwnd = ds['U10']
vwnd = ds['V10']
cosalpha = ds['COSALPHA']
sinalpha = ds['SINALPHA']

# properly convert to earth-relative U and V and convert m/s to kt
ue = 1.94386 *(uwnd.data[0]*cosalpha.data[0] - vwnd.data[0]*sinalpha.data[0]) 
ve = 1.94386 *(vwnd.data[0]*cosalpha.data[0] + uwnd.data[0]*sinalpha.data[0]) 

# make grid-relative U and V in knots
ukt = 1.94386 * uwnd.data[0]
vkt = 1.94386 * vwnd.data[0]

# Set up the projection of the data; if lat/lon then PlateCarree is what you want
plotcrs = ccrs.PlateCarree()

# this correction comes from
# https://github.com/SciTools/cartopy/issues/1179
rho = np.pi/180.
u_src_crs = ue / np.cos(lats*rho)
v_src_crs = ve
magnitude = np.sqrt(ue**2 + ve**2)
magn_src_crs = np.sqrt(u_src_crs**2 + v_src_crs**2)
ux = u_src_crs * magnitude / magn_src_crs
vx = v_src_crs * magnitude / magn_src_crs

# Set up the projection that this WRF model actually uses
#datacrs = ccrs.LambertConformal(central_longitude=-121.0,
#                               central_latitude=45.66481,
#                               standard_parallels=(60, 30),
#                               globe=ccrs.Globe(ellipse='WGS84',semimajor_axis=6370000.,semiminor_axis=6370000.))
# The following
#   ux,vx = datacrs.transform_vectors(plotcrs,lons,lats,ue,ve)
# should have worked, but it the U component
# you retrieve needs to be divided by cos(lat) to get the 
# direction correct, but then you mess up the magnitude.
# That's why we had to calculate ux and vx on our own above.
##

# Start the figure and create plot axes with proper projection
fig = plt.figure(1, figsize=(12.92, 9.85))
ax = plt.subplot(111, projection=plotcrs)
if urlat >= 0:
    ax.set_extent([urlon-10, urlon, urlat-5, urlat+3], plotcrs)
else:
    ax.set_extent([lllon, lllon+10, lllat-5, lllat+3], plotcrs)


# Add geopolitical boundaries for map reference
states = cfeature.NaturalEarthFeature(category="cultural", scale="50m",
                             facecolor="none",
                             name="admin_1_states_provinces_shp")
ax.add_feature(states, linewidth=0.5, edgecolor="black")
ax.coastlines('50m', linewidth=0.8)
ax.add_feature(cfeature.COASTLINE.with_scale('50m'))
ax.add_feature(cfeature.STATES.with_scale('50m'))

# Plot 500-hPa wind barbs in knots, regrid to reduce number of barbs
# Plot 500-hPa wind barbs in knots, regrid to reduce number of barbs
levs = np.arange(234,288,6)
cf = ax.contourf(lons, lats, t2m, levs, alpha = 0.4, transform=plotcrs)
#plt.colorbar(cf, orientation='horizontal', pad=0, aspect=50)
ax.barbs(lons, lats, ukt, vkt, pivot='tip',
         color='green', transform=plotcrs,label='Unrotated - wrong')
ax.barbs(lons, lats, ue, ve, pivot='tip',
         color='red', transform=plotcrs,label='Rotated to latlon - insufficient')
ax.barbs(lons, lats, ux, vx, pivot='tip',
         color='blue', transform=plotcrs,label='Rotated to latlon and then scaled - CORRECT')

# Make some nice titles for the plot (one right, one left)
ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
legend = plt.legend([handles[0],handles[1],handles[2]], [labels[0],labels[1],labels[2]])
# plot the points of the grid to illustrate where the vectors should be pointing
ax.plot(lons,lats,'bo', transform=plotcrs)
ax.gridlines()
plt.title('How to properly rotate WRF winds to Earth-relative coordinates\nand then reproject them to a new map projection.  Using Cartopy\n')

gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True,
                  linewidth=2, color='gray', alpha=0.5, linestyle='--')
gl.top_labels = False
gl.xformatter = LONGITUDE_FORMATTER
gl.yformatter = LATITUDE_FORMATTER
gl.ylabel_style = {'size': 10, 'color': 'black'}
gl.xlabel_style = {'size': 10, 'color': 'black'}


plt.savefig('cartopywinds.jpg')