"""Generate deterministic multi-sensor data for DOFA pipeline validation."""

from pathlib import Path

import numpy as np
import yaml


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


def make_images(count, channels, size, seed):
    rng = np.random.default_rng(seed)
    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
    images = np.empty((count, channels, size, size), dtype=np.float32)
    for sample in range(count):
        for channel in range(channels):
            phase = rng.uniform(0, 2 * np.pi)
            pattern = 0.45 + 0.25 * np.sin((channel + 1) * np.pi * x + phase)
            pattern += 0.2 * np.cos((sample % 4 + 1) * np.pi * y - phase)
            pattern += rng.normal(0, 0.02, (size, size))
            images[sample, channel] = np.clip(pattern, 0, 1)
    return images


def main():
    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    data = config["data"]
    root = ROOT / data["root"]
    root.mkdir(parents=True, exist_ok=True)
    for index, (name, modality) in enumerate(data["modalities"].items()):
        if modality.get("wavelength_mode") == "synthetic_uniform":
            wavelengths = modality["wavelength_start"] + np.arange(
                modality["channels"], dtype=np.float32
            ) * modality["wavelength_step"]
        else:
            wavelengths = np.asarray(modality["wavelengths"], dtype=np.float32)
        if wavelengths.shape != (modality["channels"],):
            raise ValueError(
                f"{name} has {modality['channels']} channels but "
                f"{wavelengths.size} wavelengths"
            )
        for split, count in (
            ("train", data["train_samples_per_modality"]),
            ("test", data["test_samples_per_modality"]),
        ):
            images = make_images(
                count,
                modality["channels"],
                data["image_size"],
                config["seed"] + index * 10 + (split == "test"),
            )
            output = root / f"{split}_{name}.npz"
            np.savez_compressed(
                output,
                images=images,
                wavelengths=wavelengths,
                modality=np.asarray(name),
                data_source=np.asarray("synthetic"),
                protocol=np.asarray(data["protocol"]),
                data_range=np.asarray(data["data_range"], dtype=np.float32),
                wavelength_mode=np.asarray(modality.get("wavelength_mode", "configured")),
            )
            print(
                f"generated={output.relative_to(ROOT)} shape={images.shape} "
                f"modality={name} channels={modality['channels']} data_source=synthetic"
            )


if __name__ == "__main__":
    main()