import os
import h5py
import numpy as np
from onescience.utils.YParams import YParams


# FGN 使用两个先验状态，自回归预测配置指定的后续状态数。
def get_dims(cfg_model, cfg_data):
    H, W = map(int, cfg_model.grid_shape)
    if tuple(map(int, cfg_data.dataset.img_size)) != (H, W):
        raise ValueError("model.grid_shape and datapipe.dataset.img_size must match")
    input_steps = int(cfg_model.input_steps)
    output_steps = int(cfg_model.output_steps)
    samples = int(cfg_data.dataloader.batch_size)
    T = input_steps + output_steps + samples - 1
    return {
        "T": T, "H": H, "W": W, "time_step": 6,
        "input_steps": input_steps, "output_steps": output_steps,
    }


def generate_fake_h5(data_dir, var_names, years, dims):
    """
    为每个年份生成一个空 h5 文件。
    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，
    文件实际只含元数据，极小，但 shape 与真实数据完全一致。
    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。

    注意：ERA5Datapipe 要求 samples_per_year = T - input_steps - output_steps + 1 >= 1，
        T 由 input_steps、output_steps 与 batch_size 自动计算。
    """
    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
    T, C = dims["T"], len(var_names)
    H, W = dims["H"], dims["W"]

    means = np.zeros((1, C, 1, 1), dtype=np.float32)
    stds = np.ones((1, C, 1, 1), dtype=np.float32)

    for year in years:
        path = os.path.join(data_dir, "data", f"{year}.h5")
        with h5py.File(path, "w") as f:
            ds = f.create_dataset(
                "fields",
                shape=(T, C, H, W),
                dtype="float32",
                chunks=(1, C, H, W),
                fillvalue=0.0,
            )
            ds.attrs["variables"] = var_names
            ds.attrs["time_step"] = dims["time_step"]
            f.create_dataset("global_means", data=means)
            f.create_dataset("global_stds", data=stds)

        size_kb = os.path.getsize(path) / 1024
        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
              f"logical={T*C*H*W*4/1024**3:.1f}GB  actual={size_kb:.1f}KB")


if __name__ == "__main__":
    cfg_model = YParams("conf/config.yaml", "model")
    cfg_datapipe = YParams("conf/config.yaml", "datapipe")

    if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"):
        print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。")
        exit()

    years = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time
    atm_vars = cfg_datapipe.dataset.channels
    if len(atm_vars) != int(cfg_model.in_channels) or len(atm_vars) != int(cfg_model.out_channels):
        raise ValueError("channel count must match model input/output channels")

    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, get_dims(cfg_model, cfg_datapipe))

    print("\n✅ Fake datasets generated.")