#!/usr/bin/env python3
"""Generate a compact deterministic ERA5-style dataset for workflow checks."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import h5py
import numpy as np
import yaml


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
    parser.add_argument("--output-dir", help="Override data.data_dir from the config")
    parser.add_argument("--height", type=int)
    parser.add_argument("--width", type=int)
    parser.add_argument("--timesteps", type=int)
    parser.add_argument("--seed", type=int)
    return parser.parse_args()


def _resolve(path: str, root: Path) -> Path:
    candidate = Path(path).expanduser()
    return candidate if candidate.is_absolute() else (root / candidate).resolve()


def generate_dataset(
    output_dir: str | Path,
    years: list[int],
    channels: list[str],
    height: int,
    width: int,
    timesteps: int,
    time_step_hours: int = 6,
    seed: int = 42,
) -> Path:
    """Write yearly HDF5 files and static fields without external data access."""
    if min(height, width, timesteps) <= 0:
        raise ValueError("height, width and timesteps must be positive")
    output_dir = Path(output_dir)
    data_dir = output_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(seed)
    channels = list(channels)
    channel_count = len(channels)
    channel_mean = np.linspace(0.0, 0.9, channel_count, dtype=np.float32)
    channel_std = np.linspace(0.8, 1.2, channel_count, dtype=np.float32)
    latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[None, :, None]
    longitude = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, None, :]
    base_pattern = 0.15 * latitude + 0.08 * np.sin(longitude)

    for year_offset, year in enumerate(sorted(set(int(item) for item in years))):
        path = data_dir / f"{year}.h5"
        with h5py.File(path, "w") as output:
            fields = output.create_dataset(
                "fields",
                shape=(timesteps, channel_count, height, width),
                dtype="float32",
                chunks=(1, channel_count, height, width),
                compression="lzf",
            )
            fields.attrs["variables"] = np.asarray(channels, dtype="S")
            fields.attrs["time_step"] = int(time_step_hours)
            for timestep in range(timesteps):
                values = np.empty((channel_count, height, width), dtype=np.float32)
                for channel in range(channel_count):
                    phase = 0.05 * (timestep + year_offset) + 0.2 * channel
                    smooth = base_pattern + 0.03 * np.sin(phase)
                    noise = rng.normal(0.0, 0.01, size=(height, width)).astype(np.float32)
                    values[channel] = channel_mean[channel] + channel_std[channel] * (smooth + noise)
                fields[timestep] = values
            output.create_dataset("global_means", data=channel_mean.reshape(1, channel_count, 1, 1))
            output.create_dataset("global_stds", data=channel_std.reshape(1, channel_count, 1, 1))
        print(f"Generated {path} shape=({timesteps},{channel_count},{height},{width})")

    static_dir = output_dir / "static"
    static_dir.mkdir(parents=True, exist_ok=True)
    latitude_2d = latitude.squeeze(0)
    land_mask = np.broadcast_to((latitude_2d > 0).astype(np.float32), (height, width)).copy()
    orography = (0.2 * np.cos(longitude) + 0.05 * latitude).astype(np.float32).squeeze(0)
    sea_ice_mask = np.broadcast_to((latitude_2d < -0.55).astype(np.float32), (height, width)).copy()
    np.save(static_dir / "land_mask.npy", land_mask)
    np.save(static_dir / "orography.npy", orography)
    np.save(static_dir / "sea_ice_mask.npy", sea_ice_mask)
    metadata = {
        "format": "ERA5-HDF5-window-v1",
        "years": sorted(set(int(item) for item in years)),
        "channels": channels,
        "shape": [timesteps, channel_count, height, width],
        "time_step_hours": int(time_step_hours),
        "seed": int(seed),
    }
    (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
    print(f"Generated static fields and metadata under {output_dir}")
    return output_dir


def main() -> None:
    args = parse_args()
    with Path(args.config).open(encoding="utf-8") as source:
        config = yaml.safe_load(source)
    data_cfg = config["data"]
    fake_cfg = config.get("fake_data", {})
    data_dir = _resolve(args.output_dir or data_cfg["data_dir"], PROJECT_ROOT)
    years = list(data_cfg["train_years"]) + list(data_cfg["val_years"]) + list(data_cfg["test_years"])
    generate_dataset(
        data_dir,
        years=years,
        channels=list(data_cfg["channels"]),
        height=int(args.height or fake_cfg.get("height", 32)),
        width=int(args.width or fake_cfg.get("width", 64)),
        timesteps=int(args.timesteps or fake_cfg.get("timesteps", 12)),
        time_step_hours=int(data_cfg.get("time_step_hours", 6)),
        seed=int(args.seed if args.seed is not None else fake_cfg.get("seed", 42)),
    )


if __name__ == "__main__":
    main()