"""Generate a small, structured PPNN dataset while preserving paper dimensions."""

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["data"]
    if (data["station_count"], data["member_count"], data["variable_count"], data["lead_hours"]) != (537, 50, 18, 48):
        raise ValueError("paper station/member/variable/lead dimensions cannot be reduced")
    rng = np.random.default_rng(int(config["seed"]))
    station_index_all = np.arange(537)
    latitude = (47.0 + 8.0 * ((station_index_all * 0.61803398875) % 1)).astype(np.float32)
    longitude = (5.5 + 9.5 * ((station_index_all * 0.41421356237) % 1)).astype(np.float32)
    elevation = (5.0 + 1450.0 * ((station_index_all * 0.2718281828) % 1)).astype(np.float32)
    station_id = np.asarray([f"DWD-{i:04d}" for i in station_index_all])
    active = np.linspace(0, 536, int(data["active_station_count"]), dtype=np.int64)
    dates = np.asarray(data["dates"], dtype="datetime64[D]")
    date_id, station_id_index = np.meshgrid(np.arange(len(dates)), active, indexing="ij")
    date_id, station_id_index = date_id.ravel(), station_id_index.ravel()
    n = len(date_id)
    day_of_year = (dates[date_id] - dates[date_id].astype("datetime64[Y]")).astype(int) + 1
    seasonal = np.sin(2 * np.pi * (day_of_year - 172) / 365.25).astype(np.float32)
    lat = latitude[station_id_index]
    lon = longitude[station_id_index]
    elev = elevation[station_id_index]
    spatial = (-0.42 * (lat - 51.0) + 0.13 * (lon - 10.0) - 0.0062 * elev).astype(np.float32)
    common_weather = rng.normal(0, 2.2, len(dates)).astype(np.float32)[date_id]
    target = (11.0 + 9.0 * seasonal + spatial + common_weather + 0.18 * np.sin(station_id_index * 0.17)).astype(np.float32)
    variable_offset = np.linspace(-4.0, 5.0, 18, dtype=np.float32)
    variable_scale = np.linspace(0.72, 1.25, 18, dtype=np.float32)
    centre = target[:, None] * variable_scale[None] + variable_offset[None] + 0.3 * common_weather[:, None]
    spread = (0.45 + 0.025 * np.abs(common_weather) + 0.00045 * elev + 0.08 * (1.0 + seasonal)).astype(np.float32)
    member_axis = np.linspace(-2.1, 2.1, 50, dtype=np.float32)
    wave = np.sin((station_id_index[:, None, None] + 1) * 0.13 + np.arange(18)[None, None] * 0.31 + member_axis[None, :, None] * 1.7)
    ensemble = centre[:, None, :] + spread[:, None, None] * variable_scale[None, None, :] * (member_axis[None, :, None] + 0.22 * wave)
    ensemble[:, :, 0] += (0.35 + 0.14 * seasonal)[:, None]
    auxiliary = np.stack((latitude[station_id_index], longitude[station_id_index], elevation[station_id_index], day_of_year / 366.0), axis=1).astype(np.float32)
    output = ROOT / data["file"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        output, format_version=np.asarray(data["format_version"]), ensemble=ensemble.astype(np.float32),
        auxiliary=auxiliary, target=target, station_index=station_id_index.astype(np.int64), date_index=date_id.astype(np.int64),
        dates=dates.astype("U10"), lead_hours=np.asarray(data["lead_hours"], np.int64), variables=np.asarray(["t2m"] + [f"var_{i:02d}" for i in range(1, 18)]),
        station_id=station_id, station_latitude=latitude, station_longitude=longitude, station_elevation_m=elevation,
    )
    if ensemble.shape != (n, 50, 18) or len(station_id) != 537 or not np.isfinite(ensemble).all():
        raise RuntimeError("generated data failed shape or finite-value validation")
    print(f"saved={output.relative_to(ROOT)} samples={n} ensemble={ensemble.shape} metadata_stations=537 lead=48h")


if __name__ == "__main__":
    main()