import sys
import pandas as pd

path = sys.argv[1] if len(sys.argv) > 1 else "all_results.csv"
df = pd.read_csv(path)
df = df[df["gt_has_gt"] == True]

for dataset in ["droid", "isaac", "libero"]:
    subset = df[df["image_path"].str.contains(f"{dataset}")]
    if subset.empty:
        print(f"{dataset}: no data")
        continue
    mae = (subset["reward"] - subset["gt_reward"]).abs().mean()
    print(f"{dataset}: MAE = {mae:.4f}  (n={len(subset)})")

avg_mae = (df["reward"] - df["gt_reward"]).abs().mean()
print(f"average: MAE = {avg_mae:.4f}  (n={len(df)})")