#!/usr/bin/env python
# coding: utf-8

# ### Demo of trying Vorticity_Divergence_Inversion.ipynb via MyBinder.org-served session
# 
# How this was generated in response [discouse conversation about running a notebook](https://discourse.jupyter.org/t/unable-to-show-the-print-output-of-my-notebook/25523/5?u=fomightez).
# (NOTE THAT BECAUSE I RAN INTO THINGS I DIDN'T UNDERSTAND I ADDED A PLOT USING A DEMO AT INPUT CELL #5 and stopped. Maybe the OP will have domain knowledge to get past this apparent hurdle that is probably due to different versions of the software being used. Ideally, the deveoper of [the source repo](https://github.com/winash12/vortdivinversion) should have listed versions in the repo with `requiremetns.txt` or `environment.yml` or in notebook with output of `%pip list` or `%conda list`, `%pip freeze` or something along those lines.)
# 
# Went to [here](https://github.com/pydata/xarray?tab=readme-ov-file#xarray-n-d-labeled-arrays-and-datasets) and clicked 'launch binder'.
# 
# In the session that came up, I opened a terminal and ran `git clone https://github.com/winash12/vortdivinversion.git` to clone the related notebook and content.
# 
# In the file navigation panel on the left, I double clicked on the directory that was listed as `vortdivinversion`.
# 
# In the top of that notebook, I made a new cell and pasted in the following code that will install the necessary packages to run the first few cells and ran it.

# In[1]:


get_ipython().run_line_magic('pip', 'install cartopy')
get_ipython().run_line_magic('pip', 'install boto3')
get_ipython().run_line_magic('pip', 'install metpy')


# In[2]:


import os.path
import xarray as xr
import numpy as np
import boto3
import metpy.calc as mpcalc
from botocore import UNSIGNED
from botocore.config import Config
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
#from matplotlib.colormaps import get_cmap
import matplotlib.ticker as mticker
import cartopy.crs as crs
from cartopy.feature import NaturalEarthFeature
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import sys
import time 
if (not os.path.isfile('gfs.t12z.pgrb2.0p25.f000')):
    client = boto3.client('s3', config=Config(signature_version=UNSIGNED))
    client.download_file('noaa-gfs-bdp-pds', 'gfs.20230809/12/atmos/gfs.t12z.pgrb2.0p25.f000', 'gfs.t12z.pgrb2.0p25.f000')


# In[3]:


u850 = xr.open_dataset('gfs.t12z.pgrb2.0p25.f000', engine='cfgrib',backend_kwargs={'filter_by_keys':{'typeOfLevel': 'isobaricInhPa', 'shortName': 'u', 'level': 850}})
u = u850.u
print(u.shape)
v850 = xr.open_dataset('gfs.t12z.pgrb2.0p25.f000', engine='cfgrib', backend_kwargs={'filter_by_keys':{'typeOfLevel': 'isobaricInhPa', 'shortName': 'v', 'level': 850}})
v = v850.v


# Compute the 850 hPa relative vorticity.

# In[4]:


vort850 = mpcalc.vorticity(u, v)
fig = plt.figure(figsize=(12,9), dpi=300.)
# Create a set of axes for the figure and set
# its map projection to that of the input data.
ax = plt.axes(projection=crs.PlateCarree())


# I don't know why that doesn't work to show any plot. (I do note that I see `UserWarning: More than one time coordinate present for variable  "u"` and a lot time warnings are just warnings and something usually shows.   
# The vorticity dpcumentation [here](https://unidata.github.io/MetPy/latest/api/generated/metpy.calc.vorticity.html) does say, "Changed in version 1.0: Changed signature from (u, v, dx, dy)". I cannot tell if maybe that is the issue and the notebook was using older version of metpy?  
# I note that the [example here](https://unidata.github.io/MetPy/latest/examples/calculations/Vorticity.html#sphx-glr-examples-calculations-vorticity-py) works fine **and gives a plot**, as can be seen by running the code below:

# In[5]:


import matplotlib.pyplot as plt

import metpy.calc as mpcalc
from metpy.cbook import example_data

# load example data
ds = example_data()

# Calculate the vertical vorticity of the flow
vort = mpcalc.vorticity(ds.uwind, ds.vwind)

# start figure and set axis
fig, ax = plt.subplots(figsize=(5, 5))

# scale vorticity by 1e5 for plotting purposes
cf = ax.contourf(ds.lon, ds.lat, vort * 1e5, range(-80, 81, 1), cmap=plt.cm.PuOr_r)
plt.colorbar(cf, pad=0, aspect=50)
ax.barbs(ds.lon.values, ds.lat.values, ds.uwind, ds.vwind, color='black', length=5, alpha=0.5)
ax.set(xlim=(260, 270), ylim=(30, 40))
ax.set_title('Relative Vorticity Calculation')

plt.show()


# Add country borders and coastlines.

# In[ ]:


countries = NaturalEarthFeature(category="cultural", scale="50m",
                                      facecolor="none",
                                      name="admin_0_countries")
ax.add_feature(countries, linewidth=.5, edgecolor="black")
ax.coastlines('50m', linewidth=0.8)


# In[ ]:


plot = vort850.plot(levels=np.arange(-1.e-4, 1.e-4, 0.2e-5), cmap=get_cmap('PRGn'), transform=crs.PlateCarree(), cbar_kwargs={'label':'relative vorticity (x$10^{-5} s^{-1}$)', 'shrink': 0.98})


# Set the map's extent to cover just Hurricane Dora.

# In[ ]:


ax.set_extent([-180.,-150.,0.,20.],crs=crs.PlateCarree())


# Add latitude/longitude gridlines.

# In[ ]:


gridlines = ax.gridlines(color="grey", linestyle="dotted", draw_labels=True)
gridlines.xlabels_top = False
gridlines.ylabels_right = False
gridlines.xlocator = mticker.FixedLocator(np.arange(-180.,149.,5.))
gridlines.ylocator = mticker.FixedLocator(np.arange(0.,21.,5.))
gridlines.xlabel_style = {'size':12, 'color':'black'}
gridlines.ylabel_style = {'size':12, 'color':'black'}
gridlines.xformatter = LONGITUDE_FORMATTER
gridlines.yformatter = LATITUDE_FORMATTER


# Add a plot title, then show the image.

# In[ ]:


plt.title("GFS 0-h 850 hPa relative vorticity (x$10^{-5} s^{-1}$) at 1200 UTC 9 August 2023")
plt.show()


# Compute the 850 hPa divergence.

# In[ ]:


div850 = mpcalc.divergence(u, v)


# Create a figure instance.

# In[ ]:


fig = plt.figure(figsize=(12,9), dpi=300.)


# Create a set of axes for the figure and set<br>
# its map projection to that of the input data.

# In[ ]:


ax = plt.axes(projection=crs.PlateCarree())


# Add country borders and coastlines.

# In[ ]:


countries = NaturalEarthFeature(category="cultural", scale="50m",
                                      facecolor="none",
                                      name="admin_0_countries")
ax.add_feature(countries, linewidth=.5, edgecolor="black")
ax.coastlines('50m', linewidth=0.8)


# Plot the 850 hPa divergence using xarray's plot functionality.

# In[ ]:


plot = div850.plot(levels=np.arange(-1.e-4, 1.e-4, 0.2e-5), cmap=get_cmap('PRGn'), transform=crs.PlateCarree(), cbar_kwargs={'label':'relative vorticity (x$10^{-5} s^{-1}$)', 'shrink': 0.98})


# Set the map's extent to cover just Hurricane Dora.

# In[ ]:


ax.set_extent([-180.,-150.,0.,20.],crs=crs.PlateCarree())


# Add latitude/longitude gridlines.

# In[ ]:


gridlines = ax.gridlines(color="grey", linestyle="dotted", draw_labels=True)
gridlines.xlabels_top = False
gridlines.ylabels_right = False
gridlines.xlocator = mticker.FixedLocator(np.arange(-180.,149.,5.))
gridlines.ylocator = mticker.FixedLocator(np.arange(0.,21.,5.))
gridlines.xlabel_style = {'size':12, 'color':'black'}
gridlines.ylabel_style = {'size':12, 'color':'black'}
gridlines.xformatter = LONGITUDE_FORMATTER
gridlines.yformatter = LATITUDE_FORMATTER


# Add a plot title, then show the image.

# In[ ]:


plt.title("GFS 0-h 850 hPa divergence (x$10^{-5} s^{-1}$) at 1200 UTC 9 August 2023")
plt.show()


# In[ ]:


vortmask = mpcalc.bounding_box_mask(vort850,5.,13.5,191.,202.)


# In[ ]:


divmask = mpcalc.bounding_box_mask(div850,5.,13.5,191.,202.)


# In[ ]:


i_bb_indices = mpcalc.find_bounding_box_indices(vortmask,5.,13.5,191.,202.)


# In[ ]:


o_bb_indices = mpcalc.find_bounding_box_indices(vortmask,0.,30,180.,220)


# In[ ]:


dx, dy = mpcalc.lat_lon_grid_deltas(vortmask.longitude, vortmask.latitude)


# In[ ]:


upsi,vpsi  = mpcalc.rotational_wind_from_inversion(vortmask,dx,dy,o_bb_indices,i_bb_indices)


# Create a figure instance.

# In[ ]:


fig = plt.figure(figsize=(12,9), dpi=300.)


# Create a set of axes for the figure and set<br>
# its map projection to that of the input data.

# In[ ]:


ax = plt.axes(projection=crs.PlateCarree())


# Add country borders and coastlines.

# In[ ]:


countries = NaturalEarthFeature(category="cultural", scale="50m",
                                      facecolor="none",
                                      name="admin_0_countries")
ax.add_feature(countries, linewidth=.5, edgecolor="black")
ax.coastlines('50m', linewidth=0.8)


# Compute the magnitude of the non-divergent component of the 850 hPa wind.

# In[ ]:


nd_spd = np.sqrt(upsi**2 + vpsi**2)


# Plot this using xarray's plot functionality.

# In[ ]:


plot = nd_spd.plot(levels=np.arange(0., 13., 1.), cmap=get_cmap('YlGnBu'), transform=crs.PlateCarree(), cbar_kwargs={'label':'non-divergent wind ($m s^{-1}$)', 'shrink': 0.98})


# Set the map's extent to match that over which we computed the non-divergent wind.

# In[ ]:


ax.set_extent([-180.,-140.,0.,30.],crs=crs.PlateCarree())


# Add latitude/longitude gridlines.

# In[ ]:


gridlines = ax.gridlines(color="grey", linestyle="dotted", draw_labels=True)
gridlines.xlabels_top = False
gridlines.ylabels_right = False
gridlines.xlocator = mticker.FixedLocator(np.arange(-180.,139.,5.))
gridlines.ylocator = mticker.FixedLocator(np.arange(0.,31.,5.))
gridlines.xlabel_style = {'size':12, 'color':'black'}
gridlines.ylabel_style = {'size':12, 'color':'black'}
gridlines.xformatter = LONGITUDE_FORMATTER
gridlines.yformatter = LATITUDE_FORMATTER


# Add a plot title, then show the image.

# In[ ]:


plt.title("GFS 0-h 850 hPa non-divergent wind magnitude ($m s^{-1}$) due to Dora at 1200 UTC 9 August 2023")
plt.show()


# In[ ]:


uchi,vchi = mpcalc.divergent_wind_from_inversion(divmask,dx,dy,o_bb_indices,i_bb_indices)


# Create a set of axes for the figure and set<br>
# its map projection to that of the input data.

# In[ ]:


ax = plt.axes(projection=crs.PlateCarree())


# Add country borders and coastlines.

# In[ ]:


countries = NaturalEarthFeature(category="cultural", scale="50m",
                                      facecolor="none",
                                      name="admin_0_countries")
ax.add_feature(countries, linewidth=.5, edgecolor="black")
ax.coastlines('50m', linewidth=0.8)


# Compute the magnitude of the non-divergent component of the 850 hPa wind.

# In[ ]:


nd_spd = np.sqrt(uchi**2 + vchi**2)


# Plot this using xarray's plot functionality.

# In[ ]:


plot = nd_spd.plot(levels=np.arange(0., 13., 1.), cmap=get_cmap('YlGnBu'), transform=crs.PlateCarree(), cbar_kwargs={'label':'non-divergent wind ($m s^{-1}$)', 'shrink': 0.98})


# Set the map's extent to match that over which we computed the non-divergent wind.

# In[ ]:


ax.set_extent([-180.,-140.,0.,30.],crs=crs.PlateCarree())


# Add latitude/longitude gridlines.

# In[ ]:


gridlines = ax.gridlines(color="grey", linestyle="dotted", draw_labels=True)
gridlines.top_labels = False
gridlines.right_labels = False
gridlines.xlocator = mticker.FixedLocator(np.arange(-180.,139.,5.))
gridlines.ylocator = mticker.FixedLocator(np.arange(0.,31.,5.))
gridlines.xlabel_style = {'size':12, 'color':'black'}
gridlines.ylabel_style = {'size':12, 'color':'black'}
gridlines.xformatter = LONGITUDE_FORMATTER
gridlines.yformatter = LATITUDE_FORMATTER


# Add a plot title, then show the image.

# In[ ]:


plt.title("GFS 0-h 850 hPa divergent wind magnitude ($m s^{-1}$) due to Dora at 1200 UTC 9 August 2023")
plt.savefig('vectorized_version')
plt.show()