import pandas as pd
import wandb

api = wandb.Api()
entity, project = "wandb_username", "kge_ogbl-biokg"
runs = api.runs(entity + "/" + project)

data = {
    "name": [],
    "NLL": [],
    "MRR": [],
    "MR": [],
    "HITS@1": [],
    "HITS@3": [],
    "HITS@10": [],
    "params": [],
}
for run in runs:
    print(run.name)
    try:
        # first, check if the run is finished
        if run.state != "finished":
            continue
        data["MRR"].append(run.summary["test/mrr"])
        data["name"].append(run.name)
        data["NLL"].append(run.summary["test/test_likelihood"])
        data["MR"].append(run.summary["test/mr"])
        data["HITS@1"].append(run.summary["test/hits@1"])
        data["HITS@3"].append(run.summary["test/hits@3"])
        data["HITS@10"].append(run.summary["test/hits@10"])
        data["params"].append(run.summary["num_parameters"])
    except KeyError:
        print(f"KeyError for run {run.name}, skipping")
        continue
df = pd.DataFrame(data)
df.to_csv(f"{project}.csv", index=False)
