from siphon.catalog import TDSCatalog
import xarray as xr
import pyvista as pv
import pvxarray
import numpy as np
import ipywidgets as widgets
import matplotlib.pyplot as plt

pv.set_plot_theme("document")
pv.set_jupyter_backend("server")  # critical for large data

# Temporary manual mode override; waiting for pyvista release including changes from https://github.com/pyvista/pyvista/pull/5798

pv.global_theme.trame.server_proxy_enabled = False
pv.global_theme.trame.jupyter_extension_enabled = True

catUrl = "https://thredds.ucar.edu/thredds/catalog/grib/NCEP/GFS/Global_0p25deg_ana/catalog.xml"
datasetName = "Full Collection Dataset";

catalog = TDSCatalog(catUrl)
ds = catalog.datasets[datasetName].remote_access(use_xarray=True)
ds

da = ds["Temperature_isobaric"]
da

dt = da[dict(time=0)]

mesh = dt.pyvista.mesh(x="lon", y="lat", z="isobaric")
mesh

pl = pv.Plotter()
pl.add_mesh(mesh, cmap="coolwarm")
pl.set_scale(zscale=0.001)
pl.show(cpos="xy")

pl = pv.Plotter()
pl.add_mesh(
    mesh.contour(
        [
            273.15,
        ]
    ),
    color=True,
)
pl.set_scale(zscale=0.001)
pl.show(cpos="xy")

pl = pv.Plotter()
pl.add_mesh(mesh.contour(), cmap="coolwarm")
pl.set_scale(zscale=0.001)
pl.show(cpos="xy")