"""Generate deterministic RemoteCLIP-format image-text pairs."""

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)
    data = config["data"]
    size = data["image_size"]
    classes = data["num_classes"]
    images = np.empty((count, 3, size, size), dtype=np.float32)
    tokens = np.zeros((count, data["context_length"]), dtype=np.int64)
    labels = np.arange(count, dtype=np.int64) % classes
    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
    for index, label in enumerate(labels):
        image = np.zeros((3, size, size), dtype=np.float32)
        image[label % 3] = 0.55 + 0.35 * np.sin((label + 1) * np.pi * x)
        image[(label + 1) % 3] += 0.25 * np.cos((label + 1) * np.pi * y)
        images[index] = np.clip(image + rng.normal(0, 0.02, image.shape), 0, 1)
        tokens[index, :4] = [label + 1, 16 + label, 32 + label, 48 + label]
    return images, tokens, labels


def main():
    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    train = make_split(config["data"]["train_samples"], config, config["seed"])
    test = make_split(config["data"]["test_samples"], config, config["seed"] + 1)
    output = ROOT / config["data"]["path"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        output,
        train_images=train[0],
        train_tokens=train[1],
        train_labels=train[2],
        test_images=test[0],
        test_tokens=test[1],
        test_labels=test[2],
        data_source=np.asarray("synthetic"),
        protocol=np.asarray(config["data"]["protocol"]),
    )
    print(
        f"generated={output.relative_to(ROOT)} train={len(train[0])} test={len(test[0])} "
        f"data_source=synthetic protocol={config['data']['protocol']}"
    )


if __name__ == "__main__":
    main()