"""Generate aligned 264x264 TerraMesh-like multimodal samples."""

from pathlib import Path

import numpy as np
import yaml


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


def make_split(path, count, config, seed):
    rng = np.random.default_rng(seed)
    size = int(config["data"]["source_size"])
    y, x = np.mgrid[-1:1:complex(size), -1:1:complex(size)].astype(np.float32)
    payload = {"format_version": np.asarray(config["data"]["format_version"]),
               "data_source": np.asarray("synthetic_terramesh_like")}
    latent = np.empty((count, size, size), np.float32)
    for sample in range(count):
        phase = rng.uniform(0, 2 * np.pi)
        latent[sample] = np.sin(3 * np.pi * x + phase) * np.cos(2 * np.pi * y - phase) + 0.3 * x + 0.2 * y
    for index, (name, channels) in enumerate(config["data"]["pixel_modalities"].items()):
        values = np.empty((count, int(channels), size, size), np.float32)
        for channel in range(int(channels)):
            values[:, channel] = latent + 0.08 * channel + 0.12 * index + rng.normal(0, 0.03, latent.shape)
        payload[f"pixel_{name}"] = values
    normalized = (latent - latent.min(axis=(1, 2), keepdims=True))
    normalized /= normalized.max(axis=(1, 2), keepdims=True).clip(1e-6)
    payload["token_map_lulc"] = np.floor(normalized * 8).clip(0, 8).astype(np.int64)
    payload["coords"] = rng.integers(0, int(config["model"]["engineering_vocab_size"]), (count, 2), dtype=np.int64)
    caption = np.empty((count, 16), np.int64)
    for sample in range(count):
        summary = int(normalized[sample].mean() * 127)
        caption[sample] = (summary + np.arange(16) * 7) % int(config["model"]["engineering_vocab_size"])
    payload["caption"] = caption
    np.savez_compressed(path, **payload)


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 offset, (filename, count) in enumerate((("train.npz", config["data"]["train_samples"]),
                                                 ("test.npz", config["data"]["test_samples"]))):
        target = output / filename
        if not target.exists():
            make_split(target, int(count), config, int(config["seed"]) + offset)
        print(f"generated={target.relative_to(ROOT)} samples={count} source_size={config['data']['source_size']}")


if __name__ == "__main__":
    main()