"""Create temporary fMoW-style temporal tensors and labels."""
import json
from pathlib import Path
import numpy as np
import yaml

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

def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    d = config["data"]
    out = ROOT / d["root"]
    out.mkdir(exist_ok=True)
    rng = np.random.default_rng(config["seed"])
    def make_split(samples):
        shape = (samples, d["frames"], d["channels"], d["image_size"], d["image_size"])
        images = rng.random(shape, dtype=np.float32)
        timestamps = np.stack(
            (
                rng.integers(0, 21, size=(samples, d["frames"])),
                rng.integers(0, 12, size=(samples, d["frames"])),
                rng.integers(0, 24, size=(samples, d["frames"])),
            ),
            axis=-1,
        ).astype(np.float32)
        order = np.argsort(timestamps[..., 0] * 12 * 24 + timestamps[..., 1] * 24 + timestamps[..., 2], axis=1)
        images = np.take_along_axis(images, order[:, :, None, None, None], axis=1)
        timestamps = np.take_along_axis(timestamps, order[..., None], axis=1)
        labels = rng.integers(d["num_classes"], size=samples, dtype=np.int64)
        return images, timestamps, labels

    train = make_split(d["train_samples"])
    test = make_split(d["test_samples"])
    np.savez_compressed(out / "train.npz", images=train[0], timestamps=train[1], labels=train[2])
    np.savez_compressed(out / "test.npz", images=test[0], timestamps=test[1], labels=test[2])
    (out / "format.json").write_text(json.dumps({
        "format": "BTCHW",
        "timestamp_format": "BT3: year_offset_2002, month_zero_based, hour",
        "source_protocol": d["protocol"],
        "data_source": "synthetic",
    }, indent=2) + "\n")
    print("created", out / "train.npz", out / "test.npz")

if __name__ == "__main__":
    main()