"""Generate compact structured daily precipitation data on the real target grid."""

from datetime import date, timedelta
from pathlib import Path

import numpy as np
import yaml


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


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    data_config = config["data"]
    samples = int(data_config["samples"])
    height, width = map(int, data_config["target_grid"])
    coarse_height, coarse_width = map(int, data_config["coarse_grid"])
    if (height, width) != (216, 488) or samples < 3:
        raise ValueError("target grid must be 216x488 and at least three days are required")
    rng = np.random.default_rng(int(config["seed"]))
    y, x = np.mgrid[-1:1:complex(height), -1:1:complex(width)]
    elevation = (2600 * np.exp(-((x + 0.2) ** 2 / 0.12 + (y - 0.05) ** 2 / 0.35)) +
                 900 * np.exp(-((x - 0.55) ** 2 / 0.035 + (y + 0.35) ** 2 / 0.09)))
    elevation += 180 * (np.sin(8 * x) * np.cos(5 * y) + 1)
    elevation = np.maximum(elevation, 0).astype(np.float32)
    target = np.empty((samples, 1, height, width), dtype=np.float32)
    for day in range(samples):
        phase = 2 * np.pi * day / max(samples, 4)
        cx, cy = -0.75 + 1.5 * day / max(samples - 1, 1), 0.35 * np.sin(phase)
        moving_system = 18 * np.exp(-((x - cx) ** 2 / 0.10 + (y - cy) ** 2 / 0.16))
        front = 7 * np.exp(-((y - 0.28 * np.sin(2 * x + phase)) ** 2) / 0.025)
        terrain_enhancement = 5.5 * (elevation / max(float(elevation.max()), 1.0)) ** 1.4
        extreme = np.zeros_like(x)
        if day % 4 == 2:
            extreme = 42 * np.exp(-((x - 0.35) ** 2 + (y + 0.18) ** 2) / 0.008)
        noise = rng.gamma(1.2, 0.35, (height, width))
        target[day, 0] = np.maximum(moving_system + front + terrain_enhancement + extreme + noise - 3.0, 0)
    # Area means produce a coarse observation; model code performs the external bilinear upsampling.
    row_edges = np.linspace(0, height, coarse_height + 1, dtype=int)
    col_edges = np.linspace(0, width, coarse_width + 1, dtype=int)
    coarse = np.empty((samples, 1, coarse_height, coarse_width), dtype=np.float32)
    for row in range(coarse_height):
        for column in range(coarse_width):
            block = target[:, :, row_edges[row]:row_edges[row + 1], col_edges[column]:col_edges[column + 1]]
            coarse[:, :, row, column] = block.mean(axis=(2, 3))
    start = date.fromisoformat(data_config["start_date"])
    timestamps = np.asarray([(start + timedelta(days=i)).isoformat() for i in range(samples)])
    years = np.asarray([(start + timedelta(days=i)).year for i in range(samples)], dtype=np.int32)
    output = ROOT / data_config["root"]
    output.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output / "daily_precipitation.npz",
                        format_version=np.array(data_config["format_version"]),
                        coarse_precipitation=coarse, elevation=elevation[None, None],
                        target_precipitation=target, timestamps=timestamps, years=years,
                        units=np.array("mm/day"), target_grid=np.array([height, width], np.int32))
    print(f"generated={samples} target_shape={target.shape} coarse_shape={coarse.shape}")


if __name__ == "__main__":
    main()