import xarray as xr
import numpy as np
import pandas as pd

# Define dimensions
lat = np.arange(-90, 91, 5)  # Latitude from -90 to 90 with 5-degree steps
lon = np.arange(-180, 181, 5) # Longitude from -180 to 180 with 5-degree steps
times = pd.date_range("2024-01-01", periods=10, freq="D") # 10 daily timesteps

# Create dummy temperature data
# Shape: (time, lat, lon)
temperature_data = np.random.rand(len(times), len(lat), len(lon)) * 30 + 273.15 # Kelvin, roughly 0-30 deg C

# Add a simple time evolution pattern (e.g., warmer over time in some regions)
# This makes the animation more interesting
for i, t in enumerate(times):
    temperature_data[i, :, :] += np.sin(np.deg2rad(lat[:, None] + i * 2)) * 5

# Create xarray DataArray
temperature_da = xr.DataArray(
    temperature_data,
    coords={"time": times, "lat": lat, "lon": lon},
    dims=["time", "lat", "lon"],
    name="temperature",
    attrs={
        "units": "K",
        "long_name": "Air Temperature",
        "standard_name": "air_temperature",
    },
)

# Create xarray Dataset
ds = xr.Dataset(
    {"temperature": temperature_da},
    attrs={"title": "Sample Temperature Data"},
)

# Save to NetCDF
output_file = "sample_temperature.nc"
ds.to_netcdf(output_file)

print(f"Sample NetCDF file '{output_file}' created successfully.")
print(ds)