import numpy as np
from pathlib import Path

def load_example_input(path: str):
    p = Path(path)
    if p.suffix == ".npz":
        z = np.load(p)
        return {k: z[k] for k in z.files}
    if p.suffix == ".npy":
        obj = np.load(p, allow_pickle=True)
        if isinstance(obj, np.ndarray) and obj.shape == ():
            return obj.item()
        if isinstance(obj, dict):
            return obj
        raise ValueError(f"Unsupported .npy content at {path}")
    raise ValueError(f"Unsupported input file type: {path}")

def save_prediction_npz(path: str, u, v, w, k):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    np.savez(path, u=u, v=v, w=w, k=k)