from datetime import datetime

import cartopy.crs as ccrs
import cartopy.feature as cfeature
from netCDF4 import Dataset, num2date
import numpy as np
from scipy.ndimage import gaussian_filter
from siphon.catalog import TDSCatalog
from siphon.ncss import NCSS
import matplotlib.pyplot as plt
import metpy.calc as mpcalc
from metpy.plots import StationPlot
from metpy.units import units

"""
Explanation: <a name="top"></a>
<div style="width:1000 px">

<div style="float:right; width:98 px; height:98px;">
<img src="https://raw.githubusercontent.com/Unidata/MetPy/master/src/metpy/plots/_static/unidata_150x150.png" alt="Unidata Logo" style="height: 98px;">
</div>

<h1>MetPy Case Study</h1>

<div style="clear:both"></div>
</div>

<hr style="height:2px;">

This is a tutorial on building a case study map for Dynamic Meteorology courses with use of Unidata tools, specifically MetPy and Siphon. In this tutorial we will cover accessing, calculating, and plotting model output.
Let's investigate The Storm of the Century, although it would easy to change which case you wanted (please feel free to do so).
Reanalysis Output: NARR 00 UTC 13 March 1993
Data from Reanalysis on pressure surfaces:

Geopotential Heights
Temperature
u-wind component
v-wind component

Calculations:

Vertical Vorticity
Advection of Temperature and Vorticity
Horizontal Divergence
Wind Speed
End of explanation
"""


# Case Study Date
year = 1993
month = 3
day = 13
hour = 0

dt = datetime(year, month, day, hour)

# Read NARR Data from THREDDS server
base_url = 'https://www.ncei.noaa.gov/thredds/catalog/narr-a-files/'

# Programmatically generate the URL to the day of data we want
cat = TDSCatalog(f'{base_url}{dt:%Y%m}/{dt:%Y%m%d}/catalog.xml')

# Have Siphon find the appropriate dataset
ds = cat.datasets.filter_time_nearest(dt)

# Download data using the NetCDF Subset Service
ncss = ds.subset()
query = ncss.query().lonlat_box(north=60, south=18, east=300, west=225)
query.all_times().variables('Geopotential_height_isobaric', 'Temperature_isobaric',
                            'u-component_of_wind_isobaric',
                            'v-component_of_wind_isobaric').add_lonlat().accept('netcdf')
data = ncss.get_data(query)

# Back up in case of bad internet connection.
# Uncomment the following line to read local netCDF file of NARR data
# data = Dataset('../../data/NARR_19930313_0000.nc','r')

"""
Explanation: Case Study Data
There are a number of different sites that you can utilize to access past model output analyses and even forecasts. The most robust collection is housed at the National Center for Environmental Information (NCEI, formerly NCDC) on a THREDDS server. The general website to begin your search is
https://www.ncdc.noaa.gov/data-access
this link contains links to many different data sources (some of which we will come back to later in this tutorial). But for now, lets investigate what model output is avaiable
https://www.ncdc.noaa.gov/data-access/model-data/model-datasets
The gridded model output that are available
Reanalysis
 * Climate Forecast System Reanalysis (CFSR)
   * CFSR provides a global reanalysis (a best estimate of the observed state of the atmosphere) of past weather from January 1979 through March 2011 at a horizontal resolution of 0.5°.
 * North American Regional Reanalysis (NARR)
   * NARR is a regional reanalysis of North America containing temperatures, winds, moisture, soil data, and dozens of other parameters at 32km horizontal resolution.
 * Reanalysis-1 / Reanalysis-2 (R1/R2)
   * Reanalysis-1 / Reanalysis-2 are two global reanalyses of atmospheric data spanning 1948/1979 to present at a 2.5° horizontal resolution.
Numerical Weather Prediction
 * Climate Forecast System (CFS) 
   * CFS provides a global reanalysis, a global reforecast of past weather, and an operational, seasonal forecast of weather out to nine months.
 * Global Data Assimilation System (GDAS)
   * GDAS is the set of assimilation data, both input and output, in various formats for the Global Forecast System model.
 * Global Ensemble Forecast System (GEFS)
   * GEFS is a global-coverage weather forecast model made up of 21 separate forecasts, or ensemble members, used to quantify the amount of uncertainty in a forecast. GEFS produces output four times a day with weather forecasts going out to 16 days.
 * Global Forecast System (GFS)
   * The GFS model is a coupled weather forecast model, composed of four separate models which work together to provide an accurate picture of weather conditions. GFS covers the entire globe down to a horizontal resolution of 28km.
 * North American Mesoscale (NAM)
   * NAM is a regional weather forecast model covering North America down to a horizontal resolution of 12km. Dozens of weather parameters are available from the NAM grids, from temperature and precipitation to lightning and turbulent kinetic energy.
 * Rapid Refresh (RAP)
   * RAP is a regional weather forecast model of North America, with separate sub-grids (with different horizontal resolutions) within the overall North America domain. RAP produces forecasts every hour with forecast lengths going out 18 hours.  RAP replaced the Rapid Update Cycle (RUC) model on May 1, 2012.
 * Navy Operational Global Atmospheric Prediction System (NOGAPS)
   * NOGAPS analysis data are available in six-hourly increments on regularly spaced latitude-longitude grids at 1-degree and one-half-degree resolutions. Vertical resolution varies from 18 to 28 pressure levels, 34 sea level depths, the surface, and other various levels.
Ocean Models
 * Hybrid Coordinate Ocean Model (HYCOM), Global
   * The Navy implementation of HYCOM is the successor to Global NCOM. This site hosts regions covering U.S. coastal waters as well as a global surface model.
 * Navy Coastal Ocean Model (NCOM), Global
   * Global NCOM was run by the Naval Oceanographic Office (NAVOCEANO) as the Navy’s operational global ocean-prediction system prior to its replacement by the Global HYCOM system in 2013. This site hosts regions covering U.S., European, West Pacific, and Australian coastal waters as well as a global surface model.
 * Navy Coastal Ocean Model (NCOM), Regional
   * The Regional NCOM is a high-resolution version of NCOM for specific areas. NCEI serves the Americas Seas, U.S. East, and Alaska regions of NCOM.
 * Naval Research Laboratory Adaptive Ecosystem Climatology (AEC)
   * The Naval Research Laboratory AEC combines an ocean model with Earth observations to provide a synoptic view of the typical (climatic) state of the ocean for every day of the year. This dataset covers the Gulf of Mexico and nearby areas.
 * National Centers for Environmental Prediction (NCEP) Real Time Ocean Forecast System (RTOFS)–Atlantic
   * RTOFS–Atlantic is a data-assimilating nowcast-forecast system operated by NCEP. This dataset covers the Gulf of Mexico and most of the northern and central Atlantic.
Climate Prediction
 * CM2 Global Coupled Climate Models (CM2.X)
   * CM2.X consists of two climate models to model the changes in climate over the past century and into the 21st century.
 * Coupled Model Intercomparison Project Phase 5 (CMIP5) (link is external)
   * The U.N. Intergovernmental Panel on Climate Change (IPCC) coordinates global analysis of climate models under the Climate Model Intercomparison Project (CMIP). CMIP5 is in its fifth iteration. Data are available through the Program for Climate Model Diagnosis and Intercomparison (PCMDI) website.
Derived / Other Model Data
 * Service Records Retention System (SRRS)
   * SRRS is a store of weather observations, summaries, forecasts, warnings, and advisories generated by the National Weather Service for public use.
 * NOMADS Ensemble Probability Tool
   * The NOMADS Ensemble Probability Tool allows a user to query the Global Ensemble Forecast System (GEFS) to determine the probability that a set of forecast conditions will occur at a given location using all of the 21 separate GEFS ensemble members.
 * National Digital Forecast Database (NDFD)
   * NDFD are gridded forecasts created from weather data collected by National Weather Service field offices and processed through the National Centers for Environmental Prediction. NDFD data are available by WMO header or by date range.
 * National Digital Guidance Database (NDGD)
   * NDGD consists of forecasts, observations, model probabilities, climatological normals, and other digital data that complement the National Digital Forecast Database.
NARR Output
Lets investigate what specific NARR output is available to work with from NCEI.
https://www.ncdc.noaa.gov/data-access/model-data/model-datasets/north-american-regional-reanalysis-narr
We specifically want to look for data that has "TDS" data access, since that is short for a THREDDS server data access point. There are a total of four different GFS datasets that we could potentially use.
Choosing our data source
Let's go ahead and use the NARR Analysis data to investigate the past case we identified (The Storm of the Century).
https://www.ncei.noaa.gov/thredds/catalog/narr-a-files/199303/19930313/catalog.html?dataset=narr-a-files/199303/19930313/narr-a_221_19930313_0000_000.grb
And we will use a python package called Siphon to read this data through the NetCDFSubset (NetCDFServer) link.
https://www.ncei.noaa.gov/thredds/ncss/grid/narr-a-files/199303/19930313/narr-a_221_19930313_0000_000.grb/dataset.html
End of explanation
"""


data.dimensions

"""
Explanation: Let's see what dimensions are in the file:
End of explanation
"""


# Extract data and assign units
tmpk = gaussian_filter(data.variables['Temperature_isobaric'][0], sigma=1.0) * units.K
hght = 0
uwnd = 0
vwnd = 0

# Extract coordinate data for plotting
lat = data.variables['lat'][:]
lon = data.variables['lon'][:]
lev = 0

"""
Explanation: Pulling Data for Calculation/Plotting
The object that we get from Siphon is netCDF-like, so we can pull data using familiar calls for all of the variables that are desired for calculations and plotting purposes.
NOTE:
Due to the curvilinear nature of the NARR grid, there is a need to smooth the data that we import for calculation and plotting purposes. For more information about why, please see the following link: http://www.atmos.albany.edu/facstaff/rmctc/narr/
Additionally, we want to attach units to our values for use in MetPy calculations later and it will also allow for easy conversion to other units.
<div class="alert alert-success">
    <b>EXERCISE</b>:
    Replace the `0`'s in the template below with your code:
     <ul>
      <li>Use the `gaussian_filter` function to smooth the `Temperature_isobaric`, `Geopotential_height_isobaric`, `u-component_of_wind_isobaric`, and `v-component_of_wind_isobaric` variables from the netCDF object with a `sigma` value of 1.</li>
    <li>Assign the units of `kelvin`, `meter`, `m/s`, and `m/s` resectively.</li>
    <li>Extract the `lat`, `lon`, and `isobaric1` variables.</li>
    </ul>
</div>
End of explanation
"""


time = data.variables['time1']
print(time.units)
vtime = num2date(time[0], units=time.units)
print(vtime)

"""
Explanation: <button data-toggle="collapse" data-target="#sol1" class='btn btn-primary'>View Solution</button>
<div id="sol1" class="collapse">
<code><pre>
# Extract data and assign units
tmpk = gaussian_filter(data.variables['Temperature_isobaric'][0],
                       sigma=1.0) * units.K
hght = gaussian_filter(data.variables['Geopotential_height_isobaric'][0],
                       sigma=1.0) * units.meter
uwnd = gaussian_filter(data.variables['u-component_of_wind_isobaric'][0], sigma=1.0) * units('m/s')
vwnd = gaussian_filter(data.variables['v-component_of_wind_isobaric'][0], sigma=1.0) * units('m/s')

\# Extract coordinate data for plotting
lat = data.variables['lat'][:]
lon = data.variables['lon'][:]
lev = data.variables['isobaric1'][:]
</pre></code>
</div>

Next we need to extract the time variable. It's not in very useful units, but the num2date function can be used to easily create regular datetime objects.
End of explanation
"""


# Calcualte dx and dy for calculations
dx, dy = mpcalc.lat_lon_grid_deltas(lon, lat)

"""
Explanation: Finally, we need to calculate the spacing of the grid in distance units instead of degrees using the MetPy helper function lat_lon_grid_spacing.
End of explanation
"""


# Specify 850 hPa data
ilev850 = np.where(lev==850)[0][0]
hght_850 = hght[ilev850]
tmpk_850 = 0
uwnd_850 = 0
vwnd_850 = 0

# Specify 500 hPa data
ilev500 = 0
hght_500 = 0
uwnd_500 = 0
vwnd_500 = 0

# Specify 300 hPa data
ilev300 = 0
hght_300 = 0
uwnd_300 = 0
vwnd_300 = 0

"""
Explanation: Finding Pressure Level Data
A robust way to parse the data for a certain pressure level is to find the index value using the np.where function. Since the NARR pressure data ('levels') is in hPa, then we'll want to search that array for our pressure levels 850, 500, and 300 hPa.
<div class="alert alert-success">
    <b>EXERCISE</b>:
    Replace the `0`'s in the template below with your code:
     <ul>
      <li>Find the index of the 850 hPa, 500 hPa, and 300 hPa levels.</li>
      <li>Extract the heights, temperature, u, and v winds at those levels.</li>
    </ul>
</div>
End of explanation
"""


# Temperature Advection
# tmpc_adv_850 = mpcalc.advection(--Fill in this call--).to('degC/s')

"""
Explanation: <button data-toggle="collapse" data-target="#sol2" class='btn btn-primary'>View Solution</button>
<div id="sol2" class="collapse">
<code><pre>
# Specify 850 hPa data
ilev850 = np.where(lev == 850)[0][0]
hght_850 = hght[ilev850]
tmpk_850 = tmpk[ilev850]
uwnd_850 = uwnd[ilev850]
vwnd_850 = vwnd[ilev850]

\# Specify 500 hPa data
ilev500 = np.where(lev == 500)[0][0]
hght_500 = hght[ilev500]
uwnd_500 = uwnd[ilev500]
vwnd_500 = vwnd[ilev500]

\# Specify 300 hPa data
ilev300 = np.where(lev == 300)[0][0]
hght_300 = hght[ilev300]
uwnd_300 = uwnd[ilev300]
vwnd_300 = vwnd[ilev300]
</pre></code>
</div>

Using MetPy to Calculate Atmospheric Dynamic Quantities
MetPy has a large and growing list of functions to calculate many different atmospheric quantities. Here we want to use some classic functions to calculate wind speed, advection, planetary vorticity, relative vorticity, and divergence.

Wind Speed: mpcalc.wind_speed()
Advection: mpcalc.advection()
Planetary Vorticity: mpcalc.coriolis_parameter()
Relative Vorticity: mpcalc.vorticity()
Divergence: mpcalc.divergence()

Note: For the above, MetPy Calculation module is imported in the following manner import metpy.calc as mpcalc.
Temperature Advection
A classic QG forcing term is 850-hPa temperature advection. MetPy has a function for advection
advection(scalar quantity, [advecting vector components], (grid spacing components))
So for temperature advection our scalar quantity would be the tempertaure, the advecting vector components would be our u and v components of the wind, and the grid spacing would be our dx and dy we computed in an earier cell.
<div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Uncomment and fill out the advection calculation below.</li>
    </ul>
</div>
End of explanation
"""


# Vorticity and Absolute Vorticity Calculations

# Planetary Vorticity
# f = mpcalc.coriolis_parameter(-- Fill in here --).to('1/s')

# Relative Vorticity
# vor_500 = mpcalc.vorticity(-- Fill in here --)

# Abosolute Vorticity
# avor_500 = vor_500 + f

"""
Explanation: <button data-toggle="collapse" data-target="#sol3" class='btn btn-primary'>View Solution</button>
<div id="sol3" class="collapse">
<code><pre>
# Temperature Advection
tmpc_adv_850 = mpcalc.advection(tmpk_850, [uwnd_850, vwnd_850],
                                (dx, dy), dim_order='yx').to('degC/s')

</pre></code>
</div>

Vorticity Calculations
There are a couple of different vorticities that we are interested in for various calculations, planetary vorticity, relative vorticity, and absolute vorticity. Currently MetPy has two of the three as functions within the calc module.
Planetary Vorticity (Coriolis Parameter)
coriolis_parameter(latitude in radians)
Note: You must can convert your array of latitudes to radians...NumPy give a great function np.deg2rad() or have units attached to your latitudes in order for MetPy to convert them for you! Always check your output to make sure that your code is producing what you think it is producing.
Relative Vorticity
When atmospheric scientists talk about relative vorticity, we are really refering to the relative vorticity that is occuring about the vertical axis (the k-hat component). So in MetPy the function is
vorticity(uwind, vwind, dx, dy)
 Absolute Vorticity
Currently there is no specific function for Absolute Vorticity, but this is easy for us to calculate from the previous two calculations because we just need to add them together!
ABS Vort = Rel. Vort + Coriolis Parameter
Here having units are great, becase we won't be able to add things together that don't have the same units! Its a nice safety check just in case you entered something wrong in another part of the calculation, you'll get a units error.
<div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Fill in the function calls below to complete the vorticity calculations.</li>
    </ul>
</div>
End of explanation
"""


# Vorticity Advection
f_adv = mpcalc.advection(f, [uwnd_500, vwnd_500], (dx, dy), dim_order='yx')

relvort_adv = mpcalc.advection(vor_500, [uwnd_500, vwnd_500], (dx, dy), dim_order='yx')

absvort_adv = mpcalc.advection(avor_500, [uwnd_500, vwnd_500], (dx, dy), dim_order='yx')

"""
Explanation: <button data-toggle="collapse" data-target="#sol4" class='btn btn-primary'>View Solution</button>
<div id="sol4" class="collapse">
<code><pre>
# Vorticity and Absolute Vorticity Calculations

\# Planetary Vorticity
f = mpcalc.coriolis_parameter(np.deg2rad(lat)).to('1/s')

\# Relative Vorticity
vor_500 = mpcalc.vorticity(uwnd_500, vwnd_500, dx, dy,
                           dim_order='yx')

\# Abosolute Vorticity
avor_500 = vor_500 + f
</pre></code>
</div>

Vorticity Advection
We use the same MetPy function for temperature advection for our vorticity advection, we just have to change the scalar quantity (what is being advected) and have appropriate vector quantities for the level our scalar is from. So for vorticity advections well want our wind components from 500 hPa.
End of explanation
"""


# Stretching Vorticity
div_500 = mpcalc.divergence(uwnd_500, vwnd_500, dx, dy, dim_order='yx')

stretch_vort = -1 * avor_500 * div_500

"""
Explanation: Divergence and Stretching Vorticity
If we want to analyze another component of the vorticity tendency equation other than advection, we might want to assess the stretching forticity term.
-(Abs. Vort.)*(Divergence)
We already have absolute vorticity calculated, so now we need to calculate the divergence of the level, which MetPy has a function
divergence(uwnd, vwnd, dx, dy)
This function computes the horizontal divergence.
End of explanation
"""


# Divergence 300 hPa, Ageostrophic Wind
wspd_300 = mpcalc.wind_speed(uwnd_300, vwnd_300).to('kts')

div_300 = mpcalc.divergence(uwnd_300, vwnd_300, dx, dy, dim_order='yx')
ugeo_300, vgeo_300 = mpcalc.geostrophic_wind(hght_300, f, dx, dy, dim_order='yx')

uageo_300 = uwnd_300 - ugeo_300
vageo_300 = vwnd_300 - vgeo_300

"""
Explanation: Wind Speed, Geostrophic and Ageostrophic Wind
Wind Speed
Calculating wind speed is not a difficult calculation, but MetPy offers a function to calculate it easily keeping units so that it is easy to convert units for plotting purposes.
wind_speed(uwnd, vwnd)
Geostrophic Wind
The geostrophic wind can be computed from a given height gradient and coriolis parameter
geostrophic_wind(heights, coriolis parameter, dx, dy)
This function will return the two geostrophic wind components in a tuple. On the left hand side you'll be able to put two variables to save them off separately, if desired.
Ageostrophic Wind
Currently, there is not a function in MetPy for calculating the ageostrophic wind, however, it is again a simple arithmatic operation to get it from the total wind (which comes from our data input) and out calculated geostrophic wind from above.
Ageo Wind = Total Wind - Geo Wind
End of explanation
"""


# Data projection; NARR Data is Earth Relative
dataproj = ccrs.PlateCarree()

# Plot projection
# The look you want for the view, LambertConformal for mid-latitude view
plotproj = ccrs.LambertConformal(central_longitude=-100., central_latitude=40.,
                                 standard_parallels=[30, 60])

def create_map_background():
    fig=plt.figure(figsize=(14, 12))
    ax=plt.subplot(111, projection=plotproj)
    ax.set_extent([-125, -73, 25, 50],ccrs.PlateCarree())
    ax.coastlines('50m', linewidth=0.75)
    ax.add_feature(cfeature.STATES, linewidth=0.5)
    return fig, ax

"""
Explanation: Maps and Projections
End of explanation
"""


fig, ax = create_map_background()

# Contour 1 - Temperature, dotted
# Your code here!

# Contour 2
clev850 = np.arange(0, 4000, 30)
cs = ax.contour(lon, lat, hght_850, clev850, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)

plt.clabel(cs, fontsize=10, inline=1, inline_spacing=10, fmt='%i',
           rightside_up=True, use_clabeltext=True)

# Filled contours - Temperature advection
contours = [-3, -2.2, -2, -1.5, -1, -0.5, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]
# Your code here!

# Vector
ax.barbs(lon, lat, uwnd_850.to('kts').m, vwnd_850.to('kts').m,
         regrid_shape=15, transform=dataproj)

# Titles
plt.title('850-hPa Geopotential Heights, Temperature (C), \
          Temp Adv (C/h), and Wind Barbs (kts)', loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()

"""
Explanation: 850-hPa Temperature Advection

Add one contour (Temperature in Celsius with a dotted linestyle
Add one colorfill (Temperature Advection in C/hr)

<div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Add one contour (Temperature in Celsius with a dotted linestyle</li>
      <li>Add one filled contour (Temperature Advection in C/hr)</li>
    </ul>
</div>
End of explanation
"""


fig, ax = create_map_background()

# Contour 1
clev500 = np.arange(0, 7000, 60)
cs = ax.contour(lon, lat, hght_500, clev500, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

# Filled contours
# Set contour intervals for Absolute Vorticity
clevavor500 = [-4, -3, -2, -1, 0, 7, 10, 13, 16, 19,
               22, 25, 28, 31, 34, 37, 40, 43, 46]

# Set colorfill colors for absolute vorticity
# purple negative
# yellow to orange positive
colorsavor500 = ('#660066', '#660099', '#6600CC', '#6600FF',
                 '#FFFFFF', '#ffE800', '#ffD800', '#ffC800',
                 '#ffB800', '#ffA800', '#ff9800', '#ff8800',
                 '#ff7800', '#ff6800', '#ff5800', '#ff5000',
                 '#ff4000', '#ff3000')

# YOUR CODE HERE!

plt.colorbar(cf, orientation='horizontal', pad=0, aspect=50)

# Vector
ax.barbs(lon, lat, uwnd_500.to('kts').m, vwnd_500.to('kts').m,
         regrid_shape=15, transform=dataproj)

# Titles
plt.title('500-hPa Geopotential Heights, Absolute Vorticity \
          (1/s), and Wind Barbs (kts)', loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()

"""
Explanation: <button data-toggle="collapse" data-target="#sol5" class='btn btn-primary'>View Solution</button>
<div id="sol5" class="collapse">
<code><pre>
fig, ax = create_map_background()

\# Contour 1 - Temperature, dotted
cs2 = ax.contour(lon, lat, tmpk_850.to('degC'), range(-50, 50, 2),
                 colors='grey', linestyles='dotted', transform=dataproj)

plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=10, fmt='%i',
           rightside_up=True, use_clabeltext=True)

\# Contour 2
clev850 = np.arange(0, 4000, 30)
cs = ax.contour(lon, lat, hght_850, clev850, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)

plt.clabel(cs, fontsize=10, inline=1, inline_spacing=10, fmt='%i',
           rightside_up=True, use_clabeltext=True)

\# Filled contours - Temperature advection
contours = [-3, -2.2, -2, -1.5, -1, -0.5, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]
cf = ax.contourf(lon, lat, tmpc_adv_850*3600, contours,
                 cmap='bwr', extend='both', transform=dataproj)
plt.colorbar(cf, orientation='horizontal', pad=0, aspect=50,
             extendrect=True, ticks=contours)

\# Vector
ax.barbs(lon, lat, uwnd_850.to('kts').m, vwnd_850.to('kts').m,
         regrid_shape=15, transform=dataproj)

\# Titles
plt.title('850-hPa Geopotential Heights, Temperature (C), \
          Temp Adv (C/h), and Wind Barbs (kts)', loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()
</pre></code>
</div>

500-hPa Absolute Vorticity
<div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Add code for plotting vorticity as filled contours with given levels and colors.</li>
    </ul>
</div>
End of explanation
"""


fig, ax = create_map_background()

# Contour 1
clev300 = np.arange(0, 11000, 120)
cs2 = ax.contour(lon, lat, div_300 * 10**5, range(-10, 11, 2),
                 colors='grey', transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

# Contour 2
cs = ax.contour(lon, lat, hght_300, clev300, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

# Filled Contours
spd300 = np.arange(50, 250, 20)
cf = ax.contourf(lon, lat, wspd_300, spd300, cmap='BuPu',
                 transform=dataproj, zorder=0)
plt.colorbar(cf, orientation='horizontal', pad=0.0, aspect=50)

# Vector of 300-hPa Ageostrophic Wind Vectors
# Your code goes here!

# Titles
plt.title('300-hPa Geopotential Heights, Divergence (1/s),\
          Wind Speed (kts), Ageostrophic Wind Vector (m/s)',
          loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()

"""
Explanation: <button data-toggle="collapse" data-target="#sol6" class='btn btn-primary'>View Solution</button>
<div id="sol6" class="collapse">
<code><pre>
fig, ax = create_map_background()

\# Contour 1
clev500 = np.arange(0, 7000, 60)
cs = ax.contour(lon, lat, hght_500, clev500, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

\# Filled contours
\# Set contour intervals for Absolute Vorticity
clevavor500 = [-4, -3, -2, -1, 0, 7, 10, 13, 16, 19,
               22, 25, 28, 31, 34, 37, 40, 43, 46]

\# Set colorfill colors for absolute vorticity
\# purple negative
\# yellow to orange positive
colorsavor500 = ('#660066', '#660099', '#6600CC', '#6600FF',
                 '#FFFFFF', '#ffE800', '#ffD800', '#ffC800',
                 '#ffB800', '#ffA800', '#ff9800', '#ff8800',
                 '#ff7800', '#ff6800', '#ff5800', '#ff5000',
                 '#ff4000', '#ff3000')

cf = ax.contourf(lon, lat, avor_500 * 10**5, clevavor500,
                 colors=colorsavor500, transform=dataproj)
plt.colorbar(cf, orientation='horizontal', pad=0, aspect=50)

\# Vector
ax.barbs(lon, lat, uwnd_500.to('kts').m, vwnd_500.to('kts').m,
         regrid_shape=15, transform=dataproj)

\# Titles
plt.title('500-hPa Geopotential Heights, Absolute Vorticity \
          (1/s), and Wind Barbs (kts)', loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()
</pre></code>
</div>

300-hPa Wind Speed, Divergence, and Ageostrophic Wind
<div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Add code to plot 300-hPa Ageostrophic Wind vectors using matplotlib's quiver function.</li>
    </ul>
</div>
End of explanation
"""


fig=plt.figure(1,figsize=(21.,16.))

# Upper-Left Panel
ax=plt.subplot(221,projection=plotproj)
ax.set_extent([-125.,-73,25.,50.],ccrs.PlateCarree())
ax.coastlines('50m', linewidth=0.75)
ax.add_feature(cfeature.STATES,linewidth=0.5)

# Contour #1
clev500 = np.arange(0,7000,60)
cs = ax.contour(lon,lat,hght_500,clev500,colors='k',
                linewidths=1.0,linestyles='solid',transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=3, fmt='%i', rightside_up=True, use_clabeltext=True)

# Contour #2
cs2 = ax.contour(lon,lat,f*10**4,np.arange(0,3,.05),colors='grey',
                linewidths=1.0,linestyles='dashed',transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=3, fmt='%.2f', rightside_up=True, use_clabeltext=True)

# Colorfill
cf = ax.contourf(lon,lat,f_adv*10**10,np.arange(-10,11,0.5),
                 cmap='PuOr_r',extend='both',transform=dataproj)
plt.colorbar(cf, orientation='horizontal',pad=0.0,aspect=50,extendrect=True)

# Vector
ax.barbs(lon,lat,uwnd_500.to('kts').m,vwnd_500.to('kts').m,regrid_shape=15,transform=dataproj)

# Titles
plt.title(r'500-hPa Geopotential Heights, Planetary Vorticity Advection ($*10^{10}$ 1/s^2)',loc='left')
plt.title('VALID: %s' %(vtime),loc='right')



# Upper-Right Panel
ax=plt.subplot(222,projection=plotproj)
ax.set_extent([-125.,-73,25.,50.],ccrs.PlateCarree())
ax.coastlines('50m', linewidth=0.75)
ax.add_feature(cfeature.STATES, linewidth=0.5)

# Contour #1
clev500 = np.arange(0,7000,60)
cs = ax.contour(lon,lat,hght_500,clev500,colors='k',
                linewidths=1.0,linestyles='solid',transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=3, fmt='%i', rightside_up=True, use_clabeltext=True)

# Contour #2
cs2 = ax.contour(lon,lat,vor_500*10**5,np.arange(-40,41,4),colors='grey',
                linewidths=1.0,transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=3, fmt='%d', rightside_up=True, use_clabeltext=True)

# Colorfill
cf = ax.contourf(lon,lat,relvort_adv*10**8,np.arange(-5,5.5,0.5),
                 cmap='BrBG',extend='both',transform=dataproj)
plt.colorbar(cf, orientation='horizontal',pad=0.0,aspect=50,extendrect=True)

# Vector
ax.barbs(lon,lat,uwnd_500.to('kts').m,vwnd_500.to('kts').m,regrid_shape=15,transform=dataproj)

# Titles
plt.title(r'500-hPa Geopotential Heights, Relative Vorticity Advection ($*10^{8}$ 1/s^2)',loc='left')
plt.title('VALID: %s' %(vtime),loc='right')



# Lower-Left Panel
ax=plt.subplot(223,projection=plotproj)
ax.set_extent([-125.,-73,25.,50.],ccrs.PlateCarree())
ax.coastlines('50m', linewidth=0.75)
ax.add_feature(cfeature.STATES, linewidth=0.5)

# Contour #1
clev500 = np.arange(0,7000,60)
cs = ax.contour(lon,lat,hght_500,clev500,colors='k',
                linewidths=1.0,linestyles='solid',transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=3, fmt='%i', rightside_up=True, use_clabeltext=True)

# Contour #2
cs2 = ax.contour(lon,lat,avor_500*10**5,np.arange(-5,41,4),colors='grey',
                linewidths=1.0,transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=3, fmt='%d', rightside_up=True, use_clabeltext=True)

# Colorfill
cf = ax.contourf(lon,lat,absvort_adv*10**8,np.arange(-5,5.5,0.5),
                 cmap='RdBu',extend='both',transform=dataproj)
plt.colorbar(cf, orientation='horizontal',pad=0.0,aspect=50,extendrect=True)

# Vector
ax.barbs(lon,lat,uwnd_500.to('kts').m,vwnd_500.to('kts').m,regrid_shape=15,transform=dataproj)

# Titles
plt.title(r'500-hPa Geopotential Heights, Absolute Vorticity Advection ($*10^{8}$ 1/s^2)',loc='left')
plt.title('VALID: %s' %(vtime),loc='right')



# Lower-Right Panel
ax=plt.subplot(224,projection=plotproj)
ax.set_extent([-125.,-73,25.,50.],ccrs.PlateCarree())
ax.coastlines('50m', linewidth=0.75)
ax.add_feature(cfeature.STATES, linewidth=0.5)

# Contour #1
clev500 = np.arange(0,7000,60)
cs = ax.contour(lon,lat,hght_500,clev500,colors='k',
                linewidths=1.0,linestyles='solid',transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=3, fmt='%i', rightside_up=True, use_clabeltext=True)

# Contour #2
cs2 = ax.contour(lon,lat,gaussian_filter(avor_500*10**5,sigma=1.0),np.arange(-5,41,4),colors='grey',
                linewidths=1.0,transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=3, fmt='%d', rightside_up=True, use_clabeltext=True)

# Colorfill
cf = ax.contourf(lon,lat,gaussian_filter(stretch_vort*10**9,sigma=1.0),np.arange(-15,16,1),
                 cmap='PRGn',extend='both',transform=dataproj)
plt.colorbar(cf, orientation='horizontal',pad=0.0,aspect=50,extendrect=True)

# Vector
ax.barbs(lon,lat,uwnd_500.to('kts').m,vwnd_500.to('kts').m,regrid_shape=15,transform=dataproj)

# Titles
plt.title(r'500-hPa Geopotential Heights, Stretching Vorticity ($*10^{9}$ 1/s^2)',loc='left')
plt.title('VALID: %s' %(vtime),loc='right')

plt.tight_layout()
plt.show()

"""
Explanation: <button data-toggle="collapse" data-target="#sol7" class='btn btn-primary'>View Solution</button>
<div id="sol7" class="collapse">
<code><pre>
fig, ax = create_map_background()

\# Contour 1
clev300 = np.arange(0, 11000, 120)
cs2 = ax.contour(lon, lat, div_300 * 10**5, range(-10, 11, 2),
                 colors='grey', transform=dataproj)
plt.clabel(cs2, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

\# Contour 2
cs = ax.contour(lon, lat, hght_300, clev300, colors='k',
                linewidths=1.0, linestyles='solid', transform=dataproj)
plt.clabel(cs, fontsize=10, inline=1, inline_spacing=4,
           fmt='%i', rightside_up=True, use_clabeltext=True)

\# Filled Contours
spd300 = np.arange(50, 250, 20)
cf = ax.contourf(lon, lat, wspd_300, spd300, cmap='BuPu',
                 transform=dataproj, zorder=0)
plt.colorbar(cf, orientation='horizontal', pad=0.0, aspect=50)

\# Vector of 300-hPa Ageostrophic Wind Vectors
ax.quiver(lon, lat, uageo_300.m, vageo_300.m, regrid_shape=15,
          pivot='mid', transform=dataproj, zorder=10)

\# Titles
plt.title('300-hPa Geopotential Heights, Divergence (1/s),\
          Wind Speed (kts), Ageostrophic Wind Vector (m/s)',
          loc='left')
plt.title(f'VALID: {vtime}', loc='right')

plt.tight_layout()
plt.show()
</pre></code>
</div>

Vorticity Tendency Terms
Here is an example of a four-panel plot for a couple of terms in the Vorticity Tendency equation
Upper-left Panel: Planetary Vorticity Advection
Upper-right Panel: Relative Vorticity Advection
Lower-left Panel: Absolute Vorticity Advection
Lower-right Panel: Stretching Vorticity
End of explanation
"""


# Set lat/lon bounds for region to plot data
LLlon = -104
LLlat = 33
URlon = -94
URlat = 38.1

# Set up mask so that you only plot what you want
skip_points = (slice(None, None, 3), slice(None, None, 3))
mask_lon = ((lon[skip_points].ravel() > LLlon + 0.05) & (lon[skip_points].ravel() < URlon + 0.01))
mask_lat = ((lat[skip_points].ravel() < URlat - 0.01) & (lat[skip_points].ravel() > LLlat - 0.01))
mask = mask_lon & mask_lat

"""
Explanation: Plotting Data for Hand Calculation
Calculating dynamic quantities with a computer is great and can allow for many different educational opportunities, but there are times when we want students to calculate those quantities by hand. So can we plot values of geopotential height, u-component of the wind, and v-component of the wind on a map? Yes! And its not too hard to do.
Since we are using NARR data, we'll plot every third point to get a roughly 1 degree by 1 degree separation of grid points and thus an average grid spacing of 111 km (not exact, but close enough for back of the envelope calculations).
To do our plotting we'll be using the functionality of MetPy to plot station plot data, but we'll use our gridded data to plot around our points. To do this we'll have to make or 2D data into 1D (which is made easy by the ravel() method associated with our data objects).
First we'll want to set some bounds (so that we only plot what we want) and create a mask to make plotting easier.
Second we'll set up our figure with a projection and then set up our "stations" at the grid points we desire using the MetPy class StationPlot
https://unidata.github.io/MetPy/latest/api/generated/metpy.plots.StationPlot.html#metpy.plots.StationPlot
Third we'll plot our points using matplotlibs scatter() function and use our stationplot object to plot data around our "stations"
End of explanation
"""


# Set up plot basics and use StationPlot class from MetPy to help with plotting
fig = plt.figure(figsize=(14, 8))
ax = plt.subplot(111,projection=ccrs.LambertConformal(central_latitude=50,central_longitude=-107))
ax.set_extent([LLlon,URlon,LLlat,URlat],ccrs.PlateCarree())
ax.coastlines('50m', edgecolor='grey', linewidth=0.75)
ax.add_feature(cfeature.STATES, edgecolor='grey', linewidth=0.5)

# Set up station plotting using only every third element from arrays for plotting
stationplot = StationPlot(ax, lon[skip_points].ravel()[mask],
                          lat[skip_points].ravel()[mask],
                          transform=ccrs.PlateCarree(), fontsize=12)

# Plot markers then data around marker for calculation purposes
# Your code goes here!

# Title
plt.title('Geopotential (m; top), U-wind (m/s; Lower Left), V-wind (m/s; Lower Right)')

plt.tight_layout()
plt.show()

"""
Explanation: <div class="alert alert-success">
    <b>EXERCISE</b>:
     <ul>
      <li>Plot markers and data around the markers.</li>
    </ul>
</div>
End of explanation
"""


