"""Create small BCHW data with the multi-scale SatMAE++ protocol."""
import json
from pathlib import Path
import numpy as np
import torch
from torch.nn import functional as F
import yaml
import argparse

ROOT = Path(__file__).resolve().parents[1]

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
    parser.add_argument("--output", type=Path, default=None)
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text())
    d, rng = cfg["data"], np.random.default_rng(cfg["seed"])
    out = args.output or ROOT / d["root"]; out.mkdir(parents=True, exist_ok=True)
    shape = (d["train_samples"], d["channels"], d["image_size"], d["image_size"])
    test_shape = (d["test_samples"], d["channels"], d["image_size"], d["image_size"])
    def make_split(split_shape):
        images = rng.random(split_shape, dtype=np.float32)
        payload = {
            "images": images,
            "labels": rng.integers(d["num_classes"], size=split_shape[0], dtype=np.int64),
        }
        tensor = torch.from_numpy(images)
        for scale in d["scales"]:
            if scale != 1:
                payload[f"images_{scale}x"] = F.interpolate(
                    tensor, scale_factor=scale, mode="bilinear", align_corners=False
                ).numpy()
        return payload

    np.savez_compressed(out / "train.npz", **make_split(shape))
    np.savez_compressed(out / "test.npz", **make_split(test_shape))
    (out / "format.json").write_text(json.dumps({
        "format": "BCHW",
        "high_resolution_fields": [f"images_{scale}x" for scale in d["scales"] if scale != 1],
        "scales": d["scales"],
        "protocol": d["protocol"],
        "data_source": "synthetic",
    }, indent=2) + "\n")
    print("created", out / "train.npz", out / "test.npz")

if __name__ == "__main__": main()