#!/usr/bin/env python3
"""
Create a simple Zarr store `foo.zarr` with monthly precip (months 1–12).
"""

import numpy as np
import xarray as xr

def main():
    months = np.arange(1, 13, dtype=np.int32)

    # Example precip values (mm/month). Replace with your real data if you have it.
    precip = np.random.rand(12).astype(np.float32) * 200.0

    ds = xr.Dataset(
        data_vars={
            "precip": (("month",), precip, {"units": "mm/month", "long_name": "Monthly precipitation"})
        },
        coords={
            "month": ("month", months, {"long_name": "Month of year"})
        },
        attrs={"title": "Example monthly precipitation", "source": "generated by script"},
    )

    # Chunk by month (nice and simple)
    ds = ds.chunk({"month": 12})

    # Write Zarr (overwrite if it exists)
    ds.to_zarr("foo.zarr", mode="w")
    print("Wrote foo.zarr")

if __name__ == "__main__":
    main()