"""Generate grouped synthetic rain days at the real 39x56x56 input shape."""

from pathlib import Path

import numpy as np
import yaml


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


def smooth(field: np.ndarray, rounds: int = 2) -> np.ndarray:
    for _ in range(rounds):
        field = sum(np.roll(np.roll(field, y, -2), x, -1) for y, x in ((0, 0), (1, 0), (-1, 0), (0, 1), (0, -1))) / 5
    return field


def make_day(rng: np.random.Generator, steps: int, height: int, width: int):
    yy, xx = np.mgrid[-1:1:complex(height), -1:1:complex(width)].astype(np.float32)
    terrain = np.clip(0.55 + 0.25 * yy + 0.16 * np.sin(2.5 * xx) - 0.12 * np.cos(3 * yy), 0, 1)
    center = rng.uniform(-0.45, 0.45, 2)
    velocity = rng.uniform(-0.10, 0.10, 2)
    width0 = rng.uniform(0.13, 0.27)
    amplitude = rng.uniform(12, 28)
    inputs = np.empty((steps, 39, height, width), np.float32)
    targets = np.empty((steps, height, width), np.float32)
    for step in range(steps):
        cy, cx = center + velocity * step
        rain = amplitude * np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * width0**2))
        rain += 0.35 * amplitude * np.exp(-((xx + cx * 0.5) ** 2 + (yy - cy * 0.6) ** 2) / (3 * width0**2))
        rain = np.clip(rain * (0.8 + 0.3 * terrain) + rng.normal(0, 0.12, rain.shape), 0, None)
        forecast = np.clip(np.roll(rain, shift=(1, -1), axis=(0, 1)) * (1.12 + 0.08 * terrain) - 0.25, 0, None)
        channels = np.empty((39, height, width), np.float32)
        for variable in range(8):
            for level in range(4):
                index = variable * 4 + level
                noise = smooth(rng.normal(0, 0.05, rain.shape).astype(np.float32))
                if variable == 1:  # humidity follows rain and terrain
                    channels[index] = 0.45 + 0.02 * rain + 0.12 * terrain - 0.025 * level + noise
                elif variable == 4:  # vertical velocity is strongest near moving rain
                    channels[index] = -0.08 * rain / (level + 1) + 0.04 * (xx * velocity[1] + yy * velocity[0]) + noise
                else:
                    channels[index] = (variable + 1) * 0.12 + level * 0.04 + 0.08 * xx - 0.05 * yy + noise
        for index in range(32, 37):
            channels[index] = 0.1 * (index - 31) + 0.03 * rain + smooth(rng.normal(0, 0.04, rain.shape).astype(np.float32))
        channels[37] = forecast
        channels[38] = terrain
        inputs[step], targets[step] = channels, rain
    return inputs, targets, terrain


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data = config["data"]
    expected = (8, 39, 56, 56)
    actual = (int(data["time_steps"]), int(data["channels"]), int(data["height"]), int(data["width"]))
    if actual != expected or len(data["lead_hours"]) != expected[0]:
        raise ValueError(f"fixed data contract is {expected} with eight lead hours, got {actual}")
    counts = data["rain_days"]
    output = ROOT / data["root"]
    output.mkdir(parents=True, exist_ok=True)
    seed = int(config["seed"])
    group_id = 0
    for split_index, split in enumerate(("train", "val", "test")):
        arrays, targets, terrains, groups = [], [], [], []
        for day in range(int(counts[split])):
            x, y, terrain = make_day(np.random.default_rng(seed + group_id), int(data["time_steps"]), int(data["height"]), int(data["width"]))
            arrays.append(x); targets.append(y); terrains.append(terrain); groups.append(group_id)
            group_id += 1
        np.savez_compressed(
            output / f"{split}.npz", inputs=np.stack(arrays), targets=np.stack(targets),
            terrain=np.stack(terrains), group_id=np.asarray(groups), split=np.asarray(split),
            format_version=np.asarray(data["format_version"]), lead_hours=np.asarray(data["lead_hours"]),
        )
        print(f"split={split} rain_days={counts[split]} inputs={np.stack(arrays).shape}")


if __name__ == "__main__":
    main()