from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
from physics import gray_scott_step, initial_state

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
SIZE = 32
TRAJECTORIES = 36
SAMPLES_PER_TRAJECTORY = 40


def main() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(2039)
    states = []
    next_states = []
    feeds = []
    kills = []
    trajectory_ids = []
    for trajectory in range(TRAJECTORIES):
        feed = float(rng.uniform(0.025, 0.060))
        kill = float(rng.uniform(0.050, 0.072))
        feed_tensor = torch.tensor([feed])
        kill_tensor = torch.tensor([kill])
        state = initial_state(1, SIZE, seed=2039 + trajectory)
        for step in range(SAMPLES_PER_TRAJECTORY * 2):
            next_state = gray_scott_step(state, feed_tensor, kill_tensor)
            if step % 2 == 0:
                states.append(state[0].numpy())
                next_states.append(next_state[0].numpy())
                feeds.append(feed)
                kills.append(kill)
                trajectory_ids.append(trajectory)
            state = next_state
    arrays = {
        "states": np.stack(states).astype(np.float32),
        "next_states": np.stack(next_states).astype(np.float32),
        "feeds": np.asarray(feeds, dtype=np.float32),
        "kills": np.asarray(kills, dtype=np.float32),
        "trajectory_ids": np.asarray(trajectory_ids, dtype=np.int64),
    }
    np.savez_compressed(DATA_DIR / "gray_scott_trajectories.npz", **arrays)
    manifest = {
        "grid_size": SIZE,
        "trajectories": TRAJECTORIES,
        "samples": len(states),
        "samples_per_trajectory": SAMPLES_PER_TRAJECTORY,
        "train_trajectories": list(range(0, 28)),
        "validation_trajectories": list(range(28, 32)),
        "test_trajectories": list(range(32, 36)),
        "feed_range": [0.025, 0.060],
        "kill_range": [0.050, 0.072],
        "path": "gray_scott_trajectories.npz",
    }
    (DATA_DIR / "manifest.json").write_text(
        json.dumps(manifest, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()