"""Generate tiny Satlas-compatible independent train/test NPZ archives."""

import json
from pathlib import Path
import numpy as np
import yaml

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


def make_split(samples, config, seed):
    rng = np.random.default_rng(seed); d = config["data"]; size = d["highres_size"]
    highres = rng.random((samples, d["highres_time_steps"], d["highres_channels"], size, size), dtype=np.float32)
    lowres = rng.random((samples, d["lowres_time_steps"], d["lowres_channels"],
                         d["lowres_size"], d["lowres_size"]), dtype=np.float32)
    valid_highres_times = np.ones((samples, d["highres_time_steps"]), dtype=np.bool_)
    valid_lowres_times = np.ones((samples, d["lowres_time_steps"]), dtype=np.bool_)
    # Exercise independent missing-observation masks even with a single smoke sample.
    valid_highres_times[:, -1] = False
    valid_lowres_times[:, -2:] = False
    highres[~valid_highres_times] = 0
    lowres[~valid_lowres_times] = 0
    yy, xx = np.mgrid[:size, :size]; segmentation = np.empty((samples, size, size), dtype=np.int64)
    point = np.zeros((samples, 1, size, size), np.float32); polygon = np.zeros_like(point); polyline = np.zeros_like(point)
    for index in range(samples):
        segmentation[index] = (xx >= size // 2) + 2 * (yy >= size // 2)
        center = size // 4 + index * 8; point[index, 0, center-2:center+3, center-2:center+3] = 1
        polygon[index, 0, size//3:2*size//3, size//3:2*size//3] = 1
        polyline[index, 0, np.arange(size), (np.arange(size) + index * 8) % size] = 1
    last_valid = valid_highres_times.sum(1) - 1
    regression = np.stack([0.7 * highres[index, last_valid[index], :1] + 0.3 * highres[index, 0, 1:2]
                           for index in range(samples)]).astype(np.float32)
    return dict(highres_images=highres, lowres_images=lowres,
                valid_highres_times=valid_highres_times, valid_lowres_times=valid_lowres_times,
                sample_ids=np.asarray([f"sample-{seed}-{index}" for index in range(samples)]), segmentation=segmentation,
                regression=regression, point=point, polygon=polygon, polyline=polyline,
                property=np.arange(samples, dtype=np.int64) % 3,
                classification=(np.arange(samples, dtype=np.int64) + 1) % 3)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()); d = config["data"]
    output = ROOT / d["root"]; output.mkdir(parents=True, exist_ok=True)
    for split, count, seed in (("train", d["train_samples"], config["seed"]), ("test", d["test_samples"], config["seed"] + 1)):
        np.savez_compressed(output / f"{split}.npz", **make_split(count, config, seed),
                             source=np.asarray("synthetic"), protocol=np.asarray(d["protocol"]))
    (output / "format.json").write_text(json.dumps({
        "protocol": d["protocol"], "source": "synthetic",
        "highres_images": "float32 [N,4,3,512,512] NAIP RGB",
        "lowres_images": "float32 [N,8,9,512,512] Sentinel-2 multispectral",
        "valid_highres_times": "bool [N,4], false slots are zero-filled",
        "valid_lowres_times": "bool [N,8], false slots are zero-filled",
        "dense_targets": "segmentation [N,512,512]; regression/point/polygon/polyline [N,1,512,512]",
        "global_targets": "property/classification int64 [N]"
    }, indent=2) + "\n")
    print(f"created {output / 'train.npz'} and {output / 'test.npz'}")


if __name__ == "__main__": main()