import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import metpy.calc as mpcalc
import metpy.constants as mpconstants
from metpy.units import units
import numpy as np
import xarray as xr
from datetime import datetime, timedelta, time 

e = datetime(2023, 5, 1, 12)
ds = xr.open_dataset(f'https://www.ncei.noaa.gov/thredds/dodsC/model-gfs-g4-anl-files/'
                     f'{date:%Y%m}/{date:%Y%m%d}/gfs_4_{date:%Y%m%d_%H}00_000.grb2')

# Sunset geographic extent of data to CONUS to limit how much is downloaded
lon_slice = slice(-125, -67)  
lat_slice = slice(24, 50)     

lats = ds['lat'].values  # Extract latitude
lons = ds['lon'].values  # Extract longitude

# Create meshgrid from 1D lat and lon arrays
lons_grid, lats_grid = np.meshgrid(lons, lats)

# Compute grid spacings (dx, dy) using np.gradient
dy, dx = np.gradient(lats_grid, axis=0), np.gradient(lons_grid, axis=1)

# Grab data and smooth using a nine-point filter applied 50 times to grab the synoptic signal
level = 700 * units.hPa
hght_700 = ds['Geopotential_height_isobaric'].metpy.sel(vertical=level)
tmpk_700 = ds['Temperature_isobaric'].metpy.sel(vertical=level)
uwnd_700 = ds['u-component_of_wind_isobaric'].metpy.sel(vertical=level)
vwnd_700 = ds['v-component_of_wind_isobaric'].metpy.sel(vertical=level)
w_700 = ds['Vertical_velocity_geometric_isobaric'].metpy.sel(vertical=level)

# Compute the Q-vector components
uqvect, vqvect = mpcalc.q_vector(uwnd_700, vwnd_700, tmpk_700, level)
print(uqvect)
print(uqvect.to_base_units())
# Compute the divergence of the Q-vectors calculated above (term A of the RHS according to Hoskins et al. (1978)
q_div = -2 * mpcalc.divergence(uqvect, vqvect)