#!/usr/bin/env python3
"""Generate MRMS-shaped synthetic events for contract and smoke checks."""
from __future__ import annotations

import argparse
import struct
import zlib
from pathlib import Path

import numpy as np


def write_png_gray16(path: Path, array: np.ndarray) -> None:
    array = np.asarray(array, dtype=">u2")
    raw = b"".join(b"\x00" + row.tobytes() for row in array)
    def chunk(kind: bytes, payload: bytes) -> bytes:
        return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", zlib.crc32(kind + payload) & 0xffffffff)
    header = struct.pack(">IIBBBBB", array.shape[1], array.shape[0], 16, 0, 0, 0, 0)
    path.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress(raw, 1)) + chunk(b"IEND", b""))


def generate(output_dir: str | Path, events: int = 2, height: int = 512, width: int = 512, total_length: int = 29, seed: int = 42) -> None:
    root = Path(output_dir)
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[:height, :width]
    for event_index in range(events):
        event = root / f"synthetic_{event_index:04d}"
        event.mkdir(parents=True, exist_ok=True)
        cx, cy = width * (0.25 + 0.2 * event_index), height * 0.45
        for frame in range(total_length):
            center_x = cx + frame * 1.5
            rain = 12.0 * np.exp(-((xx - center_x) ** 2 + (yy - cy) ** 2) / (2 * (max(height, width) * 0.12) ** 2))
            rain += rng.normal(0, 0.08, size=(height, width))
            # Inverse official encoding: decoded = uint16 / 10 - 3.
            encoded = np.clip(np.rint((np.maximum(rain, 0) + 3.0) * 10), 0, 65535).astype(np.uint16)
            write_png_gray16(event / f"{event.name}-{frame:02d}.png", encoded)
    print(f"generated {events} MRMS events at {root} with shape ({total_length},{height},{width})")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=str, default="data/data")
    parser.add_argument("--events", type=int, default=2)
    parser.add_argument("--height", type=int, default=512)
    parser.add_argument("--width", type=int, default=512)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    generate(args.output_dir, args.events, args.height, args.width, seed=args.seed)