from __future__ import annotations

import argparse
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import h5py
import numpy as np

from common import DEFAULT_CONFIG, load_config, resolve_path


def generate_year(
    path: Path,
    variables: list[str],
    *,
    time_steps: int,
    height: int,
    width: int,
    time_step_hours: int,
    chunk_time_steps: int,
    fill_value: float,
    materialize_pattern: bool,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    channels = len(variables)
    with h5py.File(path, "w") as handle:
        fields = handle.create_dataset(
            "fields",
            shape=(time_steps, channels, height, width),
            dtype=np.float32,
            chunks=(chunk_time_steps, 1, height, width),
            fillvalue=np.float32(fill_value),
            compression="lzf",
        )
        fields.attrs["variables"] = np.asarray(variables, dtype=h5py.string_dtype())
        fields.attrs["time_step"] = time_step_hours
        handle.create_dataset(
            "global_means", data=np.zeros((1, channels, 1, 1), dtype=np.float32)
        )
        handle.create_dataset(
            "global_stds", data=np.ones((1, channels, 1, 1), dtype=np.float32)
        )

        if materialize_pattern:
            latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]
            longitude = np.linspace(
                0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32
            )
            base = latitude + np.sin(longitude)[None, :]
            # Two frames are enough to exercise non-zero input and target reads.
            for time_index in range(min(time_steps, 2)):
                for channel_index in range(channels):
                    fields[time_index, channel_index] = (
                        base + channel_index / channels + time_index * 0.01
                    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate ERA5-compatible FCNv2 data")
    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--no-pattern", action="store_true")
    args = parser.parse_args()

    config = load_config(args.config)
    data = config["data"]
    fake = config["fake_data"]
    output_dir = resolve_path(config, data["dataset_dir"])
    years = sorted(set(data["train_years"] + data["val_years"] + data["test_years"]))
    height, width = data["grid_shape"]
    for year in years:
        path = output_dir / "data" / f"{year}.h5"
        generate_year(
            path,
            data["variables"],
            time_steps=fake["time_steps_per_year"],
            height=height,
            width=width,
            time_step_hours=data["time_step_hours"],
            chunk_time_steps=fake["chunk_time_steps"],
            fill_value=fake["fill_value"],
            materialize_pattern=fake["materialize_pattern"] and not args.no_pattern,
        )
        logical_gib = (
            fake["time_steps_per_year"] * len(data["variables"]) * height * width * 4
        ) / 1024**3
        print(f"{path}: logical={logical_gib:.2f} GiB, actual={path.stat().st_size / 1024**2:.2f} MiB")


if __name__ == "__main__":
    main()