from typing import Dict, Union

import wandb


class WandbLogger:
    def __init__(
        self,
        project_name: str,
        experiment_name: str,
        config: Dict[str, Union[int, str]],
    ):
        self.project_name = project_name
        self.config = config
        self.experiment_name = experiment_name
        wandb.init(project=project_name, name=experiment_name, config=config)  # type: ignore

    def log(self, data: Dict[str, Union[int, float]]) -> None:
        wandb.log(data)  # type: ignore
