"""Create synthetic temperature data for the EarthInversion NetCDF tutorials.

Run: python write_netcdf.py [example.nc]
These values are teaching data, not observations or a physical model.
"""
from datetime import datetime, timedelta
from pathlib import Path

import numpy as np
from netCDF4 import Dataset, date2num


def write_example(path="example.nc"):
    latitudes = np.linspace(35.0, 39.0, 5, dtype="f4")
    longitudes = np.linspace(-124.0, -119.0, 6, dtype="f4")
    dates = [datetime(2026, 1, 1) + timedelta(hours=6 * i) for i in range(3)]
    temperature = (
        280.0 + 0.5 * np.arange(3)[:, None, None]
        + 0.2 * (latitudes[None, :, None] - 35.0)
        + 0.1 * (longitudes[None, None, :] + 124.0)
    ).astype("f4")
    temperature = np.ma.array(temperature, mask=False)
    temperature[0, 0, 0] = np.ma.masked

    # Refuse to overwrite an existing file.
    with Dataset(path, "w", format="NETCDF4", clobber=False) as ds:
        ds.createDimension("time", None)
        ds.createDimension("lat", len(latitudes))
        ds.createDimension("lon", len(longitudes))
        ds.title = "Synthetic temperature example"
        ds.source = "Deterministic teaching data; not observations or model output"

        lat = ds.createVariable("lat", "f4", ("lat",))
        lat.standard_name = "latitude"
        lat.units = "degrees_north"
        lat.axis = "Y"
        lon = ds.createVariable("lon", "f4", ("lon",))
        lon.standard_name = "longitude"
        lon.units = "degrees_east"
        lon.axis = "X"
        time = ds.createVariable("time", "f8", ("time",))
        time.standard_name = "time"
        time.units = "hours since 2026-01-01 00:00:00"
        time.calendar = "standard"
        time.axis = "T"
        temp = ds.createVariable(
            "temperature", "f4", ("time", "lat", "lon"),
            fill_value=-9999.0, compression="zlib", complevel=4,
        )
        temp.standard_name = "air_temperature"
        temp.long_name = "Synthetic air temperature"
        temp.units = "K"
        lat[:] = latitudes
        lon[:] = longitudes
        time[:] = date2num(dates, units=time.units, calendar=time.calendar)
        temp[:] = temperature
    print(f"Wrote {path}: 3 times, 5 latitudes, 6 longitudes")
    return Path(path)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("path", nargs="?", default="example.nc")
    args = parser.parse_args()
    write_example(args.path)