scripts/fake_data.py

"""
Generate synthetic ERA5 data and normalization files for Stormer pipeline.

Creates:
1. HDF5 data files (lightweight, chunked with zero fill)
2. Static/invariant files (geopotential, land_sea_mask, lat, lon, etc.)
3. Normalization constants matching official Stormer format:
   - normalize_mean.npz, normalize_std.npz (input normalization, 109+ vars)
   - normalize_diff_mean_{t}.npz, normalize_diff_std_{t}.npz (diff normalization)

Expected output structure:
    data/
    ├── data/
    │   ├── {year}.h5
    │   └── ...
    ├── static/
    │   ├── geopotential.nc, land_sea_mask.nc
    │   ├── land_mask.npy, soil_type.npy, topography.npy
    │   ├── lat.npy, lon.npy
    └── normalize/
        ├── normalize_mean.npz, normalize_std.npz
        └── normalize_diff_mean_{6,12,24}.npz, normalize_diff_std_{6,12,24}.npz
"""

import os
import sys
from pathlib import Path

root_path = Path(__file__).parent.parent
sys.path.append(str(root_path))

import h5py
import numpy as np
import xarray as xr

from onescience.utils.YParams import YParams

# Stormer data dimensions:
# T=20 gives 120hrs (5 days) per year — enough for 72h lead time validation
DATASET_DIMS = {"T": 20, "H": 128, "W": 256, "time_step": 6}

# Full variable list from WeatherBench2 (109 vars) — used for normalization files
WB2_ALL_VARS = [
    # Surface
    "2m_temperature", "10m_u_component_of_wind", "10m_v_component_of_wind",
    "mean_sea_level_pressure",
    # Geopotential at pressure levels
    "geopotential_50", "geopotential_100", "geopotential_150", "geopotential_200",
    "geopotential_250", "geopotential_300", "geopotential_400", "geopotential_500",
    "geopotential_600", "geopotential_700", "geopotential_850", "geopotential_925",
    "geopotential_1000",
    # U wind
    "u_component_of_wind_50", "u_component_of_wind_100", "u_component_of_wind_150",
    "u_component_of_wind_200", "u_component_of_wind_250", "u_component_of_wind_300",
    "u_component_of_wind_400", "u_component_of_wind_500", "u_component_of_wind_600",
    "u_component_of_wind_700", "u_component_of_wind_850", "u_component_of_wind_925",
    "u_component_of_wind_1000",
    # V wind
    "v_component_of_wind_50", "v_component_of_wind_100", "v_component_of_wind_150",
    "v_component_of_wind_200", "v_component_of_wind_250", "v_component_of_wind_300",
    "v_component_of_wind_400", "v_component_of_wind_500", "v_component_of_wind_600",
    "v_component_of_wind_700", "v_component_of_wind_850", "v_component_of_wind_925",
    "v_component_of_wind_1000",
    # Temperature
    "temperature_50", "temperature_100", "temperature_150", "temperature_200",
    "temperature_250", "temperature_300", "temperature_400", "temperature_500",
    "temperature_600", "temperature_700", "temperature_850", "temperature_925",
    "temperature_1000",
    # Specific humidity
    "specific_humidity_50", "specific_humidity_100", "specific_humidity_150",
    "specific_humidity_200", "specific_humidity_250", "specific_humidity_300",
    "specific_humidity_400", "specific_humidity_500", "specific_humidity_600",
    "specific_humidity_700", "specific_humidity_850", "specific_humidity_925",
    "specific_humidity_1000",
    # Additional WB2 variables (for normalization compatibility)
    "angle_of_sub_gridscale_orography", "anisotropy_of_sub_gridscale_orography",
    "geopotential_at_surface", "high_vegetation_cover", "lake_cover", "lake_depth",
    "land_sea_mask", "low_vegetation_cover", "orography",
    "slope_of_sub_gridscale_orography", "soil_type",
    "standard_deviation_of_filtered_subgrid_orography",
    "standard_deviation_of_orography", "type_of_high_vegetation",
    "type_of_low_vegetation",
    "mean_surface_latent_heat_flux", "mean_surface_net_long_wave_radiation_flux",
    "mean_surface_net_short_wave_radiation_flux", "mean_surface_sensible_heat_flux",
    "mean_top_downward_short_wave_radiation_flux",
    "mean_top_net_long_wave_radiation_flux", "mean_top_net_short_wave_radiation_flux",
    "skin_temperature", "snow_depth", "10m_wind_speed", "surface_pressure",
    "toa_incident_solar_radiation", "total_precipitation_6hr",
    "total_column_water_vapour", "total_cloud_cover", "sea_ice_cover",
    "sea_surface_temperature", "vertical_velocity_50", "vertical_velocity_100",
    "vertical_velocity_150", "vertical_velocity_200", "vertical_velocity_250",
    "vertical_velocity_300", "vertical_velocity_400", "vertical_velocity_500",
    "vertical_velocity_600", "vertical_velocity_700", "vertical_velocity_850",
    "vertical_velocity_925", "vertical_velocity_1000",
]


def generate_fake_h5(data_dir, var_names, years, dims):
    """Generate HDF5 files with random values for pipeline validation.

    Writes actual random data (N(0, 1) per variable) to ensure non-zero
    training loss for verifying gradient flow.
    """
    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
    T, C = dims["T"], len(var_names)
    H, W = dims["H"], dims["W"]

    for year in years:
        path = os.path.join(data_dir, "data", f"{year}.h5")

        # Generate random data with per-variable mean=0, std=1
        rng = np.random.default_rng(42 + year)
        fields = rng.normal(0, 1, (T, C, H, W)).astype(np.float32)

        # Add temporal correlation so consecutive frames are similar
        # (simple AR(1): smooth the time dimension)
        for t in range(1, T):
            fields[t] = 0.7 * fields[t-1] + 0.3 * fields[t]

        # Per-variable means/stds
        means = fields.mean(axis=(0, 2, 3), keepdims=True).reshape(1, C, 1, 1)
        stds = fields.std(axis=(0, 2, 3), keepdims=True).reshape(1, C, 1, 1)
        stds = np.maximum(stds, 0.01)  # avoid division by zero

        with h5py.File(path, "w") as f:
            ds = f.create_dataset(
                "fields",
                data=fields,
                dtype="float32",
                chunks=(1, C, H, W),
            )
            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_mb = os.path.getsize(path) / 1024 / 1024
        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
              f"size={size_mb:.1f}MB  logical={T*C*H*W*4/1024**3:.1f}GB")


def get_static(data_dir):
    """Generate synthetic static/invariant data files."""
    os.makedirs(data_dir, exist_ok=True)

    H, W = DATASET_DIMS["H"], DATASET_DIMS["W"]

    # Geopotential at surface
    ds = xr.Dataset(
        data_vars={
            "z": (("valid_time", "latitude", "longitude"),
                  np.random.rand(1, H, W).astype(np.float32))
        },
        coords={
            "valid_time": ["2015-12-31"],
            "latitude": np.linspace(90, -90, H, dtype=np.float64),
            "longitude": np.linspace(0, 359.75, W, dtype=np.float64),
            "number": 0, "expver": "",
        },
        attrs={
            "GRIB_centre": "ecmf",
            "Conventions": "CF-1.7",
            "institution": "European Centre for Medium-Range Weather Forecasts",
            "history": "Generated for Stormer onescience",
        }
    )
    ds.to_netcdf(f"{data_dir}/geopotential.nc")

    # Land-sea mask
    ds_lsm = xr.Dataset(
        data_vars={
            "lsm": (("valid_time", "latitude", "longitude"),
                    np.random.rand(1, H, W).astype(np.float32))
        },
        coords={
            "valid_time": ["2015-12-31"],
            "latitude": np.linspace(90, -90, H, dtype=np.float64),
            "longitude": np.linspace(0, 359.75, W, dtype=np.float64),
            "number": 0, "expver": "",
        },
    )
    ds_lsm.to_netcdf(f"{data_dir}/land_sea_mask.nc")

    # Static numpy arrays
    arr = np.random.randn(H, W).astype(np.float32)
    np.save(f'{data_dir}/land_mask.npy', arr)
    np.save(f'{data_dir}/soil_type.npy', arr)
    np.save(f'{data_dir}/topography.npy', arr)

    # Latitude/longitude arrays
    lat = np.linspace(90, -90, H, dtype=np.float32)
    np.save(f'{data_dir}/lat.npy', lat)
    lon = np.linspace(0, 359.75, W, dtype=np.float32)
    np.save(f'{data_dir}/lon.npy', lon)

    print(f"✅ Static data generated in {data_dir}")


def generate_normalization_files(normalize_dir, all_vars, stormer_vars):
    """Generate normalization .npz files matching official Stormer format.

    Creates files for ALL 109+ WB2 variables (for compatibility with official
    normalization loading code), plus diff normalization for intervals [6, 12, 24].

    For fake data: input mean=0, std=1; diff mean=0, diff_std scales with interval.
    """
    os.makedirs(normalize_dir, exist_ok=True)

    # ---- Input normalization ----
    # mean = 0, std = 1 for all variables (fake zero-mean unit-variance data)
    inp_mean = {v: np.array([0.0], dtype=np.float32) for v in all_vars}
    inp_std = {v: np.array([1.0], dtype=np.float32) for v in all_vars}

    np.savez(os.path.join(normalize_dir, "normalize_mean.npz"), **inp_mean)
    np.savez(os.path.join(normalize_dir, "normalize_std.npz"), **inp_std)
    print(f"  normalize_mean.npz: {len(inp_mean)} vars")
    print(f"  normalize_std.npz: {len(inp_std)} vars")

    # ---- Diff normalization for each interval ----
    # diff_std scales with sqrt(interval) (diffusion-like behavior in atmosphere)
    for interval in [6, 12, 24]:
        scale = np.sqrt(interval / 6.0)  # 6h→1.0, 12h→1.414, 24h→2.0
        diff_mean = {v: np.array([0.0], dtype=np.float32) for v in all_vars}
        diff_std = {v: np.array([scale], dtype=np.float32) for v in all_vars}

        np.savez(os.path.join(normalize_dir, f"normalize_diff_mean_{interval}.npz"),
                 **diff_mean)
        np.savez(os.path.join(normalize_dir, f"normalize_diff_std_{interval}.npz"),
                 **diff_std)
        print(f"  normalize_diff_*_{interval}.npz: {len(diff_std)} vars (scale={scale:.3f})")

    print(f"✅ Normalization files generated in {normalize_dir}")


if __name__ == "__main__":
    cfg_datapipe = YParams(os.path.join(str(root_path), "conf/config.yaml"), "datapipe")
    cfg_model = YParams(os.path.join(str(root_path), "conf/config.yaml"), "model")

    data_dir = cfg_datapipe.dataset.data_dir

    # Safety: refuse to overwrite real data paths
    if data_dir.startswith("/public/") or data_dir.startswith("/work2/") or data_dir.startswith("/work/"):
        print("❌ 请检查 config，确保 data_dir 指向本地测试路径而非生产路径。")
        sys.exit(1)

    years = (cfg_datapipe.dataset.train_time +
             cfg_datapipe.dataset.val_time +
             cfg_datapipe.dataset.test_time)
    stormer_vars = cfg_datapipe.dataset.channels

    # 1. Generate fake HDF5 data
    print(f"\n📂 Generating fake HDF5: {len(years)} years, {len(stormer_vars)} variables...")
    generate_fake_h5(data_dir, stormer_vars, years, DATASET_DIMS)

    # 2. Generate static files
    static_dir = os.path.join(data_dir, "static")
    print(f"\n📂 Generating static files...")
    get_static(static_dir)

    # 3. Generate normalization files
    normalize_dir = cfg_model.normalize_dir
    print(f"\n📂 Generating normalization files...")
    generate_normalization_files(normalize_dir, WB2_ALL_VARS, stormer_vars)

    print(f"\n✅ All fake datasets generated successfully.")
    print(f"   Data dir:    {data_dir}")
    print(f"   Norm dir:    {normalize_dir}")
    print(f"   Variables:   {len(stormer_vars)} (Stormer) / {len(WB2_ALL_VARS)} (WB2 full)")
    print(f"   Years:       {years}")
    print(f"   Resolution:  {DATASET_DIMS['H']}×{DATASET_DIMS['W']} (1.40625°)")