"""Generate virtual data in the native Samudra NPZ layout."""

try:
    from ._bootstrap import ROOT
except ImportError:
    from _bootstrap import ROOT

import argparse
from pathlib import Path

import numpy as np
import yaml

STATE_CHANNELS = 77
BOUNDARY_CHANNELS = 4


def generate(path: str | Path, time: int, height: int, width: int, seed: int) -> None:
    if time < 10:
        raise ValueError("time must be at least 10 for four-pass recurrent training")
    rng = np.random.default_rng(seed)
    prognostic = rng.standard_normal((time, STATE_CHANNELS, height, width), dtype=np.float32)
    boundary = rng.standard_normal((time, BOUNDARY_CHANNELS, height, width), dtype=np.float32)
    output = Path(path)
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, prognostic=prognostic, boundary=boundary)
    print(f"saved native data: {output} prognostic={prognostic.shape} boundary={boundary.shape}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="./conf/config.yaml")
    parser.add_argument("--train-output", default="./data/train.npz")
    parser.add_argument("--test-output", default="./data/test.npz")
    parser.add_argument("--time", type=int, default=None)
    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()
    with open(args.config, encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    fake = config.get("fake_data", {})
    time = args.time or int(fake.get("time", 12))
    height = args.height or int(fake.get("height", 32))
    width = args.width or int(fake.get("width", 64))
    seed = args.seed if args.seed is not None else int(config["project"].get("seed", 1))
    generate(args.train_output, time, height, width, seed)
    generate(args.test_output, time, height, width, seed + 1)


if __name__ == "__main__":
    main()