from pathlib import Path
import sys
import numpy as np
import torch.nn.functional as F
import torch
ROOT = Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
from model.scale_adaptive_cm import load_config, structured_fields

c = load_config(ROOT); h,w=c["data"]["high_grid"]; lh,lw=c["data"]["low_grid"]
high=structured_fields(c["data"]["samples"],h,w,c["seed"])
low=F.avg_pool2d(torch.from_numpy(high),c["data"]["scale_factor"]).numpy()
assert low.shape[-2:]==(lh,lw)
path=ROOT/c["data"]["path"]; path.parent.mkdir(parents=True,exist_ok=True)
np.savez_compressed(path,format_version=np.array(c["data"]["format_version"]),low=low,high=high,split=np.array(["train"]*(len(high)-1)+["test"]),unit=np.array("mm day-1"))
print(path,low.shape,high.shape)