"""Generate native-grid virtual ERA5 data for OneScience ERA5Dataset."""

import argparse
from pathlib import Path

import numpy as np

from config import load_config, project_path

def generate(dataset_dir, year, time_steps, height, width, variables, time_step_hours, seed, dtype):
    import h5py

    dataset_dir = Path(dataset_dir)
    data_dir = dataset_dir / "data"
    stats_dir = dataset_dir / "stats"
    data_dir.mkdir(parents=True, exist_ok=True)
    stats_dir.mkdir(parents=True, exist_ok=True)
    path = data_dir / f"{year}.h5"
    rng = np.random.default_rng(seed + year)
    lat = np.linspace(90.0, -90.0, height, dtype=np.float32)[:, None]
    lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)[None, :]
    with h5py.File(path, "w") as handle:
        fields = handle.create_dataset(
            "fields", shape=(time_steps, len(variables), height, width), dtype=dtype,
            chunks=(1, 1, min(height, 64), min(width, 128)),
        )
        fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype())
        fields.attrs["time_step"] = time_step_hours
        for step in range(time_steps):
            phase = 2.0 * np.pi * (step / max(time_steps, 1))
            base = np.sin(np.deg2rad(lat)) + 0.25 * np.cos(np.deg2rad(lon) + phase)
            noise = rng.normal(0.0, 0.01, size=(len(variables), height, width)).astype(np.float32)
            fields[step] = np.asarray([(index + 1) * base + noise[index] for index in range(len(variables))], dtype=dtype)
    # ERA5Dataset always initializes these files, even when normalize=False.
    np.save(stats_dir / "global_means.npy", np.zeros((1, len(variables), 1, 1), dtype=np.float32))
    np.save(stats_dir / "global_stds.npy", np.ones((1, len(variables), 1, 1), dtype=np.float32))
    return path


def main():
    config = load_config()
    data_config = config["data"]
    runtime_config = config["runtime"]
    inference_config = config["inference"]
    parser = argparse.ArgumentParser(description="Generate native-grid virtual ERA5 HDF5 data.")
    parser.add_argument("--dataset-dir", default=str(project_path(data_config["virtual_dir"])))
    parser.add_argument("--year", type=int, default=inference_config["year"])
    parser.add_argument("--time-steps", type=int, default=data_config["virtual_time_steps"])
    parser.add_argument("--height", type=int, default=data_config["raw_height"])
    parser.add_argument("--width", type=int, default=data_config["raw_width"])
    parser.add_argument("--seed", type=int, default=runtime_config["seed"])
    parser.add_argument("--dtype", choices=["float16", "float32"], default=data_config["virtual_dtype"])
    args = parser.parse_args()
    print(generate(args.dataset_dir, args.year, args.time_steps, args.height, args.width,
                   data_config["variables"], data_config["time_step_hours"], args.seed, args.dtype))


if __name__ == "__main__":
    main()