from pathlib import Path
import sys

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.ace2 import (
    PRECIPITATION,
    Q_INDICES,
    RADIATION_INDICES,
    SURFACE_PRESSURE,
    forcing_for_hours,
    load_config,
    seed_all,
)


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


def main():
    cfg = load_config(ROOT)
    seed_all(cfg["seed"])
    d = cfg["data"]
    shape = (d["samples"], d["time_steps"], d["channels"], d["height"], d["width"])
    rng = np.random.default_rng(cfg["seed"])
    lat = np.deg2rad(np.linspace(-89.5, 89.5, d["height"], dtype=np.float32))[None, None, None, :, None]
    lon = np.deg2rad(np.linspace(0.5, 359.5, d["width"], dtype=np.float32))[None, None, None, None, :]
    time = np.arange(d["time_steps"], dtype=np.float32)[None, :, None, None, None]
    channel = np.arange(d["channels"], dtype=np.float32)[None, None, :, None, None]
    state = (0.2 * np.sin(lat * (1 + channel % 3)) + 0.1 * np.cos(lon + time / 4)
             + 0.002 * channel + 0.003 * time).astype(np.float32)
    state = np.broadcast_to(state, shape).copy()
    state += rng.normal(0, 0.005, (d["samples"], 1, d["channels"], 1, 1)).astype(np.float32)
    state[:, :, Q_INDICES] = np.maximum(state[:, :, Q_INDICES] * 0.01 + 0.003, 0)
    state[:, :, SURFACE_PRESSURE] = 1.0 + 0.01 * np.cos(lat[:, :, 0])
    state[:, :, PRECIPITATION] = np.maximum(0, 0.001 * (np.sin(lon[:, :, 0] + time[:, :, 0]) + 1))
    state[:, :, RADIATION_INDICES] = np.maximum(state[:, :, RADIATION_INDICES] + 0.5, 0)
    hours = (np.arange(d["samples"])[:, None] * d["time_steps"] + np.arange(d["time_steps"])[None]) * d["dt_hours"]
    forcing = forcing_for_hours(hours.reshape(-1), d["height"], d["width"]).reshape(
        d["samples"], d["time_steps"], 4, d["height"], d["width"]
    )
    path = ROOT / d["path"]
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, state=state.astype(np.float16), forcing=forcing.astype(np.float16), hours=hours)
    print(f"created {path}: state={state.shape}, forcing={forcing.shape}, dt={d['dt_hours']}h")


if __name__ == "__main__":
    main()