"""Generate a tiny, structurally realistic CorrDiff NPZ dataset."""

import argparse
from pathlib import Path

import numpy as np
import yaml

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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
    parser.add_argument("--output")
    args = parser.parse_args()
    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
    data = config["data"]
    count = data["fake_samples"]
    rng = np.random.default_rng(config["seed"])
    coarse = rng.normal(size=(count, *data["input_shape"])).astype("float32")
    # Correlated targets make the regression/backward smoke test meaningful.
    y = np.linspace(-1, 1, data["target_shape"][1], dtype="float32")
    x = np.linspace(-1, 1, data["target_shape"][2], dtype="float32")
    yy, xx = np.meshgrid(y, x, indexing="ij")
    target = np.empty((count, *data["target_shape"]), dtype="float32")
    coarse_signal = coarse.mean(axis=(2, 3))
    for sample in range(count):
        for channel in range(data["target_shape"][0]):
            target[sample, channel] = coarse_signal[sample, channel] + 0.3 * np.sin(
                (channel + 1) * np.pi * xx
            ) + 0.2 * np.cos((channel + 1) * np.pi * yy)
    target += rng.normal(0, 0.05, target.shape).astype("float32")
    output = Path(args.output) if args.output else ROOT / data["path"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, input=coarse, target=target,
                        protocol=np.asarray(data["protocol"]), data_source=np.asarray("synthetic"))
    print(f"saved={output} input={coarse.shape} target={target.shape}")


if __name__ == "__main__":
    main()