"""Minimal W&B integration for a PRL-launched experiment.

PRL injects PRL_RUN_ID, PRL_TASK_ID, PRL_HYPOTHESIS_ID, PRL_GIT_COMMIT and
WANDB_* variables before executing this program.
"""
import os
import wandb

run = wandb.init(
    id=os.environ.get("WANDB_RUN_ID"),
    name=os.environ.get("WANDB_NAME"),
    group=os.environ.get("WANDB_GROUP"),
    tags=os.environ.get("WANDB_TAGS", "").split(",") if os.environ.get("WANDB_TAGS") else [],
    config={
        "prl_run_id": os.environ.get("PRL_RUN_ID"),
        "prl_task_id": os.environ.get("PRL_TASK_ID"),
        "prl_hypothesis_id": os.environ.get("PRL_HYPOTHESIS_ID"),
        "git_commit": os.environ.get("PRL_GIT_COMMIT"),
    },
)
for step in range(10):
    wandb.log({"loss": 1.0 / (step + 1), "step": step})
run.finish()