"""Generate deterministic virtual data for SEEDS pipeline validation."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import yaml


def _load_config(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def generate_dataset(path: Path, samples: int, channels: int, faces: int, height: int, width: int, seed_count: int, rng: np.random.Generator) -> None:
    climate = rng.normal(0.0, 0.5, size=(samples, channels, faces, height, width)).astype(np.float32)
    seeds = (climate[:, None] + rng.normal(0.0, 0.8, size=(samples, seed_count, channels, faces, height, width))).astype(np.float32)
    targets = (climate + 0.25 * seeds.mean(axis=1) + rng.normal(0.0, 0.5, size=climate.shape)).astype(np.float32)
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, seeds=seeds, targets=targets, climate=climate)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--samples", type=int, default=8)
    parser.add_argument("--height", type=int, default=None)
    parser.add_argument("--width", type=int, default=None)
    parser.add_argument("--seed", type=int, default=None)
    args = parser.parse_args()
    config = _load_config(args.config)
    data = config["data"]
    paths = config["paths"]
    height = args.height or data["height"]
    width = args.width or data["width"]
    rng = np.random.default_rng(args.seed if args.seed is not None else config["project"]["seed"])
    for name in ("train_data", "val_data", "test_data"):
        generate_dataset(
            Path(paths[name]), args.samples, len(data["variables"]), data["faces"], height, width, data["seed_count"], rng
        )
        print(f"generated {paths[name]}")


if __name__ == "__main__":
    main()