import argparse
from pathlib import Path
import numpy as np


def main():
    parser = argparse.ArgumentParser(description="Evaluate CorrDiff NPZ predictions")
    parser.add_argument("--prediction", default="result/output/predictions.npz")
    parser.add_argument("--target", default="data/era5_corrdiff.npz")
    args = parser.parse_args()
    prediction = np.load(args.prediction)["prediction"]
    target = np.load(args.target)["target"]
    mae = np.abs(prediction.mean(axis=0) - target).mean()
    output = Path(args.prediction).with_name("metrics.npz")
    np.savez(output, mae=np.array(mae, dtype="float32"))
    print(f"ensemble={prediction.shape[0]} mae={mae:.6f}")
    print(f"saved: {output}")


if __name__ == "__main__":
    main()