"""Create an ERA5-style HDF5 dataset consumable by OneScience ERA5Dataset."""

from __future__ import annotations

import argparse
from pathlib import Path

import h5py
import numpy as np

from common import load_config, resolve_path
from variables import c85_from_config


def create_year(
    path: Path,
    year: int,
    steps: int,
    height: int,
    width: int,
    channels: list[str],
    time_step_hours: int,
    hdf5_config: dict,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(path, "w") as handle:
        fields = handle.create_dataset(
            "fields",
            shape=(steps, len(channels), height, width),
            dtype="float32",
            chunks=tuple(min(size, limit) for size, limit in zip((steps, len(channels), height, width), hdf5_config["chunks"])),
            fillvalue=0.0,
            compression=hdf5_config["compression"],
            compression_opts=hdf5_config["compression_level"],
        )
        fields.attrs["variables"] = np.asarray(channels, dtype=h5py.string_dtype("utf-8"))
        fields.attrs["time_step"] = time_step_hours
        fields.attrs["year"] = year
        lat = np.linspace(90.0, -90.0, height, dtype=np.float32)[:, None]
        lon = np.arange(width, dtype=np.float32)[None, :] * (360.0 / width)
        t2m_index = channels.index("t2m")
        for step in range(steps):
            fields[step, t2m_index] = (
                np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon + step * 15.0))
            )
        handle.create_dataset("global_means", data=np.zeros((1, len(channels), 1, 1), dtype=np.float32))
        handle.create_dataset("global_stds", data=np.ones((1, len(channels), 1, 1), dtype=np.float32))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="conf/config.yaml")
    args = parser.parse_args()
    cfg = load_config(args.config)
    data_cfg = cfg["data"]
    channels, _ = c85_from_config(cfg)
    root = resolve_path(cfg["paths"]["data_root"], cfg)
    generated_years = set()
    for split, split_cfg in data_cfg["splits"].items():
        for year in split_cfg["years"]:
            if year in generated_years:
                raise ValueError(f"Year {year} is assigned to more than one data split")
            generated_years.add(year)
            create_year(
                root / "data" / f"{year}.h5",
                year,
                split_cfg["time_steps"],
                *data_cfg["grid_size"],
                channels,
                data_cfg["time_step_hours"],
                data_cfg["hdf5"],
            )
    print(f"Created ERA5-compatible yearly datasets at {root / 'data'}")


if __name__ == "__main__":
    main()