from __future__ import annotations

import argparse
from pathlib import Path
import sys

import h5py
import numpy as np
import yaml


PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "scripts"))

from era5_adapter import inspect_era5_contract


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def generate_year_file(
    output_path: Path,
    variables: list[str],
    time_steps: int,
    height: int,
    width: int,
    time_step_hours: int,
) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    channels = len(variables)
    means = np.zeros((1, channels, 1, 1), dtype=np.float32)
    stds = np.ones((1, channels, 1, 1), dtype=np.float32)

    with h5py.File(output_path, "w") as handle:
        fields = handle.create_dataset(
            "fields",
            shape=(time_steps, channels, height, width),
            dtype="float32",
            chunks=(1, 1, height, width),
            fillvalue=0.0,
        )
        fields.attrs["variables"] = variables
        fields.attrs["time_step"] = time_step_hours
        handle.create_dataset("global_means", data=means)
        handle.create_dataset("global_stds", data=stds)


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate lightweight ERA5-style HDF5 files for W-MAE.")
    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf" / "config.yaml")
    parser.add_argument("--output-dir", type=Path, default=None)
    args = parser.parse_args()

    config = load_config(args.config)
    data_config = config["data"]
    fake_config = config["fake_data"]
    output_dir = args.output_dir or PROJECT_ROOT / data_config["dataset_dir"]
    variables = list(fake_config["variables"])
    years = list(fake_config["years"])
    source_height, source_width = data_config["source_size"]

    for year in years:
        path = output_dir / "data" / f"{year}.h5"
        generate_year_file(
            output_path=path,
            variables=variables,
            time_steps=int(fake_config["time_steps"]),
            height=int(source_height),
            width=int(source_width),
            time_step_hours=int(data_config["time_step_hours"]),
        )
        print(f"created {path} ({path.stat().st_size / 1024:.1f} KiB physical size)")

    report = inspect_era5_contract(output_dir, years, variables)
    print(f"validated fields shape: {report['fields_shape']}")
    print(f"W-MAE spatial adapter: {report['crop']} -> {report['model_size']}")
    print("Synthetic channel names are test-only and must not be used for real training.")


if __name__ == "__main__":
    main()