from hypnotoad.cases.circular import CircularEquilibrium
from hypnotoad.core.mesh import BoutMesh
import numpy as np

R0 = 2.3
B0 = 3.2
r_inner = 0.1
r_outer = 1.4
q = 4.1
settings = {
    "R0": R0,
    "B0": B0,
    "poloidal_spacing_method": "linear",
    "q_coefficients": [q],
    "r_inner": r_inner,
    "r_outer": r_outer,
    "refine_methods": "line",
    "refine_width": 1.0e-3,
}

equilib = CircularEquilibrium(settings, nonorthogonal_settings=settings)

mesh = BoutMesh(equilib, settings)
mesh.geometry()
mesh.writeGridfile("grid.nc")

from matplotlib import pyplot as plt
from xbout import open_boutdataset
#from xbout.load import _open_grid

gridfilepath = "grid.nc"
grid = open_boutdataset(gridfilepath, geometry="toroidal")
#grid = _open_grid(gridfilepath, chunks={},keep_xboundaries=True,keep_yboundaries=True)

grid["psi_poloidal"].bout.contourf()
grid["psi_poloidal"].bout.contour()
grid["psi_poloidal"].bout.pcolormesh()
grid["psi_poloidal"].bout.pcolormesh(shading="gouraud")
grid["psi_poloidal"].bout.regions()
plt.show()

$ python3 plot_test.py
NXPE not found, setting to 1
NYPE not found, setting to 1
MXG not found, setting to 2
MYG not found, setting to 0
MXSUB not found, setting to 2
MYSUB not found, setting to 8
Traceback (most recent call last):
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1317, in _construct_dataarray
    variable = self._variables[name]
               ~~~~~~~~~~~~~~~^^^^^^
KeyError: 't_array'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "*/plot_test.py", line 11, in <module>
    grid = open_boutdataset(gridfilepath, geometry="toroidal")
  File "*/lib/python3.13/site-packages/xbout/load.py", line 279, in open_boutdataset
    ds, remove_yboundaries = _auto_open_mfboutdataset(
                             ~~~~~~~~~~~~~~~~~~~~~~~~^
        datapath=datapath,
        ^^^^^^^^^^^^^^^^^^
    ...<4 lines>...
        **kwargs,
        ^^^^^^^^^
    )
    ^
  File "*/lib/python3.13/site-packages/xbout/load.py", line 736, in _auto_open_mfboutdataset
    _, unique_indices = unique(ds["t_array"], return_index=True)
                               ~~^^^^^^^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1410, in __getitem__
    return self._construct_dataarray(key)
           ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 1319, in _construct_dataarray
    _, name, variable = _get_virtual_variable(self._variables, name, self.dims)
                        ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "*/lib/python3.13/site-packages/xarray/core/dataset.py", line 175, in _get_virtual_variable
    raise KeyError(key)
KeyError: 't_array'