"""Generate tiny paper-shape RemoteCLIP train/test NPZ files."""

import json
from pathlib import Path

import numpy as np
import yaml

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


def make_split(count, config, seed):
    rng = np.random.default_rng(seed)
    d = config["data"]
    images = np.empty((count, 3, 224, 224), dtype=np.float32)
    tokens = np.zeros((count, 77), dtype=np.int64)
    pair_ids = np.arange(count, dtype=np.int64) % d["num_semantic_groups"]
    yy, xx = np.mgrid[:224, :224].astype(np.float32) / 223
    for index, pair_id in enumerate(pair_ids):
        base = np.stack((xx, yy, (xx + yy) / 2)) if pair_id == 0 else np.stack((yy, 1 - xx, xx * yy))
        images[index] = np.clip(base + rng.normal(0, 0.025, base.shape), 0, 1)
        # CLIP convention: EOT has the largest vocabulary id and therefore wins argmax pooling.
        tokens[index, :6] = [49406, 100 + pair_id, 200 + pair_id, 300 + index, 400 + pair_id, 49407]
    return images, tokens, pair_ids


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    output = ROOT / config["data"]["root"]
    output.mkdir(parents=True, exist_ok=True)
    for split, count, seed in (("train", config["data"]["train_samples"], config["seed"]),
                               ("test", config["data"]["test_samples"], config["seed"] + 1)):
        images, tokens, pair_ids = make_split(count, config, seed)
        np.savez_compressed(output / f"{split}.npz", images=images, tokens=tokens, pair_ids=pair_ids,
                            data_source=np.asarray("synthetic"), protocol=np.asarray(config["data"]["protocol"]))
    (output / "format.json").write_text(json.dumps({
        "protocol": config["data"]["protocol"], "data_source": "synthetic",
        "images": "float32 [N,3,224,224] in [0,1]", "tokens": "int64 [N,77] CLIP BPE ids",
        "pair_ids": "int64 [N], equal ids define valid multi-positive matches"
    }, indent=2) + "\n")
    print(f"created {output / 'train.npz'} and {output / 'test.npz'}")


if __name__ == "__main__":
    main()