from pathlib import Path
import argparse
import numpy as np


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", default="data/fake_corrdiff.npz")
    parser.add_argument("--samples", type=int, default=1)
    args = parser.parse_args()
    path = Path(args.output)
    path.parent.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(42)
    inputs = rng.standard_normal((args.samples, 12, 36, 36)).astype("float32")
    targets = rng.standard_normal((args.samples, 4, 448, 448)).astype("float32")
    np.savez_compressed(path, input=inputs, target=targets)
    print(f"input: {inputs.shape}")
    print(f"target: {targets.shape}")
    print(f"saved: {path}")


if __name__ == "__main__":
    main()