try:
    import xarray
except ModuleNotFoundError:
    !sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet xarray)"
    pass

#
import numpy as np

print("numpy version=", np.__version__)

#
import xarray as xr

print("xarray version=", xr.__version__)

# from IPython.display import HTML, display

# display(HTML("<style>.container { width:77% !important; }</style>"))

np.random.seed(314)

# Create a 2D array.
data = np.array([[1, 2, 3], [4, 5, 6]])
print("data=", data)

# Create an xarray.
data = xr.DataArray(
    data,
    # - Assign x and y to the dimensions.
    dims=("x", "y"),
    # - Assign coordinate labels 10 and 20 to locations along x dimension.
    coords={"x": [10, 20]},
)

# Show as HTML.
data

print(type(data))

# Print as string.
print(data)

# Extract the numpy data structure.
vals = data.values
print("type(vals)=", type(vals))
print("vals=\n", vals)

# Extract the dimension names which are a tuple.
print(data.dims)

# Extract the coordinates.
print("# type(data.coords)=\n%s" % type(data.coords))
print("# data.coords=\n%s" % data.coords)

# Extract the attributes, which can store arbitrary metadata.
data.attrs

data

# Slice like NumPy-style.
# Set the x dimension to be 0, so get the first row.
data[0, :]

# Set the x dimension to be 1 (like numpy), so get the second row.
data[1, :]

# loc, "location": select by coordinate label (like pandas)
# Get data along the first dimension for the index called `10`.
data.loc[10]

# isel, "integer select": select by dimension name and integer label
# Get data along the dimension `x` for the first index
data.isel(x=0)

# isel, "integer select": select by dimension name and integer label
# Get data along the dimension `y` for the second index
data.isel(y=1)

# sel, "select", by dimension name and coordinate label
# Get data along the dimension `x` and the index `10`
data.sel(x=10)

data.attrs["long_name"] = "random_velocity"
data.attrs["units"] = "m/s"
data.attrs["description"] = "A random var created as an example"

print(data.attrs)

data

data.x.attrs["units"] = "x units"

data

# Sum 10 element-wise.
data + 10

# Sum elements across all the dimensions
data.sum()

# Compute mean along one dimension by label.
data.mean(dim="x")

# Transpose.
data.T

# Create a 1-vector with coordinates along the `y` axis.
a = xr.DataArray([1, 2, 3], {"y": [0, 1, 2]})

# Create a 1-vector with dimension `z` and no coordinates.
b = xr.DataArray([10, 20, 30, 40], dims="z")

display(a)
display(b)

# The broadcast happens along the dimensions by name, without worrying about the order.
a + b

# Create a 1-vector with coordinates along `z`.
a = xr.DataArray([1, 2, 3], {"z": [0, 1, 2]})
# Create a 1-vector with `z` dimension.
b = xr.DataArray([10, 20, 30], dims="z")

a + b

# The plot uses the attributes.
data.plot()

# From xarray to multi-index pd.Series.
srs = data.to_series()
srs

# From pd.Series to xarray.
srs.to_xarray()

df = data.to_dataframe(name="hello")
df

# Create a dictionary with heterogeneous data.
dict_ = dict(foo=data, bar=("x", [1, 2]), baz=np.pi)
print(dict_)

# Create a `xarray.Dataset` from the dictionary.
# - `foo` is a DataArray (with 2 dimensions `x` and `y`)
# - `bar` is a one-dimensional array with dimension `x`
# - `baz` is a scalar (with no dimensions)
ds = xr.Dataset(dict_)
ds

# Extract one variable from the dataset.
ds["foo"]
# This is equivalent to
# ds.foo
assert str(ds.foo) == str(ds["foo"])

ds["bar"]

ds["baz"]

ds[["foo", "bar"]].to_dataframe().values

# Both `foo` and `bar` variables have the same coordinate `x`, so we can use `x` to slice the data.
ds["x"]
# ds.bar["x"]
# ds.foo["x"]