"""Generate a compact, full-resolution synthetic RYDL-like sequence."""

from datetime import datetime, timedelta
from pathlib import Path
import sys

import h5py
import numpy as np
import yaml

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


def load_config():
    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def precipitation_frame(y, x, step, rng):
    cells = (
        (220 + 4.2 * step, 260 + 6.0 * step, 70 + 0.5 * step, 45, 1.25),
        (610 - 3.2 * step, 590 + 2.7 * step, 58, 82 - 0.4 * step, 0.85),
        (430 + 1.5 * step, 710 - 4.0 * step, 42 + 0.6 * step, 55, 0.55),
    )
    field = np.zeros_like(x, dtype=np.float32)
    for cy, cx, sy, sx, amplitude in cells:
        evolution = 1.0 + 0.12 * np.sin((step + amplitude) / 3.0)
        field += amplitude * evolution * np.exp(
            -0.5 * (((x - cx) / sx) ** 2 + ((y - cy) / sy) ** 2)
        )
    # A smooth perturbation evolves with the cells without decorrelating frames.
    phase = rng.uniform(-0.03, 0.03)
    field *= 1.0 + 0.025 * np.sin(x / 35.0 + step / 4.0 + phase) * np.cos(y / 47.0)
    return np.maximum(field, 0).astype(np.float32)


def main():
    config = load_config()
    data = config["data"]
    np.random.seed(config["seed"])
    rng = np.random.default_rng(config["seed"])
    total = data["train_frames"] + data["val_frames"] + data["test_frames"]
    height, width = data["raw_height"], data["raw_width"]
    yy, xx = np.mgrid[:height, :width].astype(np.float32)
    path = ROOT / data["path"]
    path.parent.mkdir(parents=True, exist_ok=True)
    start = datetime(2017, 1, 1)
    keys = []
    with h5py.File(path, "w") as handle:
        handle.attrs["units"] = "mm/5min"
        handle.attrs["interval_minutes"] = data["interval_minutes"]
        for index in range(total):
            key = (start + timedelta(minutes=index * data["interval_minutes"])).strftime(
                "%Y%m%d%H%M"
            )
            keys.append(key)
            handle.create_dataset(key, data=precipitation_frame(yy, xx, index, rng))
    with h5py.File(path, "r") as handle:
        first = handle[keys[0]][...]
    print(f"Number of frames: {len(keys)}")
    print(f"First key: {keys[0]}")
    print(f"Last key: {keys[-1]}")
    print(f"Frame shape: {first.shape}")
    print(f"dtype: {first.dtype}")
    print(f"min: {first.min():.8f}")
    print(f"max: {first.max():.8f}")
    print(f"mean: {first.mean():.8f}")


if __name__ == "__main__":
    main()