"""Write a lazy 569-station protocol manifest; fields are generated per chunk."""

from pathlib import Path
import numpy as np
import yaml

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


def station_metadata(count: int):
    index = np.arange(count, dtype=np.float32)
    latitude = -43.5 + 33.0 * ((index * 0.61803398875) % 1)
    longitude = 113.0 + 40.0 * ((index * 0.41421356237) % 1)
    elevation = 20 + 1450 * ((index * 0.2718281828) % 1)
    return latitude, longitude, elevation.astype(np.float32)


def make_split(path: Path, config: dict):
    data = config["data"]
    dates = np.asarray(data["valid_dates"], dtype="U10")
    lat, lon, elevation = station_metadata(data["station_count"])
    history_dates = np.empty((len(dates), data["history_days"]), dtype="U10")
    for i, date in enumerate(dates.astype("datetime64[D]")):
        history_dates[i] = (date - np.arange(data["history_days"], 0, -1)).astype("U10")
    np.savez(path, format_version=np.asarray(data["format_version"]), valid_dates=dates,
             history_dates=history_dates, initialization_utc=np.asarray("1200"), lead_hours=np.arange(241),
             station_id=np.asarray([f"JIVE-{i:04d}" for i in range(data["station_count"])]),
             station_latitude=lat, station_longitude=lon, station_elevation_m=elevation,
             variables=np.asarray(data["variables"]), units=np.asarray(data["units"]), sources=np.asarray(data["sources"]),
             neighborhood_shape=np.asarray([3, 3]), target_grid_projection=np.asarray("Albers"),
             target_grid_nx=np.asarray(None), target_grid_ny=np.asarray(None),
             representation=np.asarray(data["target_representation"]))


def generate_chunk(meta, date_index: int, station_selector, history: bool, seed: int, include_patch: bool = True):
    """Emulate lazy reads after authoritative station patches have been extracted."""
    station = np.arange(len(meta["station_id"]))[station_selector]
    days = 30 if history else 1
    lead = np.arange(241, dtype=np.float32)[None, :, None, None]
    day = np.arange(days, dtype=np.float32)[:, None, None, None]
    lat = meta["station_latitude"][station][None, None, None, :]
    lon = meta["station_longitude"][station][None, None, None, :]
    phase = date_index * 0.7 + day * 0.11
    temp = 20 - 0.35 * (lat + 25) + 4 * np.sin(2 * np.pi * (lead + 12) / 24 + phase)
    dew = temp - 5 - 1.5 * np.cos(np.deg2rad(lon) + lead / 48)
    wind = 5 + 1.2 * np.abs(np.sin(np.deg2rad(lon) + lead / 18 + phase))
    truth = np.concatenate((temp, dew, wind), axis=2).astype(np.float32)
    source = np.arange(3, dtype=np.float32)[None, None, None, :, None]
    forecast = truth[:, :, :, None] + (source - 0.6) * np.asarray([1.1, 0.8, 0.5], np.float32)[None, None, :, None, None]
    forecast += (lead[..., None] / 240) * np.asarray([0.8, -0.5, 0.6], np.float32)[None, None, :, None, None]
    rng = np.random.default_rng(seed + date_index * 1000 + int(station[0]))
    forecast += rng.normal(0, 0.08, forecast.shape).astype(np.float32)
    offsets = np.asarray([[-0.18, -0.10, -0.04], [-0.08, 0.0, 0.09], [0.03, 0.12, 0.20]], np.float32)
    elevation_delta = meta["station_elevation_m"][station] - np.mean(meta["station_elevation_m"])
    if not include_patch:
        return forecast.astype(np.float32), truth, elevation_delta.astype(np.float32)
    patches = forecast[..., None, None] + offsets
    return patches.astype(np.float32), truth, elevation_delta.astype(np.float32)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    output = ROOT / config["data"]["root"]
    output.mkdir(parents=True, exist_ok=True)
    make_split(output / "protocol.npz", config)
    print("generated=data/protocol.npz dates=2 history_days=30 leads=241 stations=569 patch=3x3 nx=unknown ny=unknown")


if __name__ == "__main__":
    main()